diff --git a/ksp/generator/src/main/kotlin/WalkOnKSFiles.kt b/ksp/generator/src/main/kotlin/WalkOnKSFiles.kt new file mode 100644 index 00000000000..8634cf815b7 --- /dev/null +++ b/ksp/generator/src/main/kotlin/WalkOnKSFiles.kt @@ -0,0 +1,30 @@ +package dev.inmo.micro_ksp.generator + +import com.google.devtools.ksp.getAllSuperTypes +import com.google.devtools.ksp.symbol.KSAnnotated +import com.google.devtools.ksp.symbol.KSClassDeclaration +import com.google.devtools.ksp.symbol.KSDeclarationContainer +import com.google.devtools.ksp.symbol.KSFile + +fun KSClassDeclaration.findSubClasses(subSymbol: KSAnnotated): Sequence { + return when (subSymbol) { + is KSClassDeclaration -> if (subSymbol.getAllSuperTypes().map { it.declaration }.contains(this)) { + sequenceOf(subSymbol) + } else { + sequenceOf() + } + else -> sequenceOf() + } + if (subSymbol is KSDeclarationContainer) { + subSymbol.declarations.flatMap { + findSubClasses(it) + } + } else { + sequenceOf() + } +} + +fun KSClassDeclaration.findSubClasses(files: Sequence): Sequence { + return files.flatMap { + findSubClasses(it) + } +} diff --git a/ksp/sealed/generator/src/main/kotlin/Processor.kt b/ksp/sealed/generator/src/main/kotlin/Processor.kt index e8816e7bb14..de653834d08 100644 --- a/ksp/sealed/generator/src/main/kotlin/Processor.kt +++ b/ksp/sealed/generator/src/main/kotlin/Processor.kt @@ -6,7 +6,6 @@ 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.AnnotationSpec import com.squareup.kotlinpoet.ClassName import com.squareup.kotlinpoet.CodeBlock import com.squareup.kotlinpoet.FileSpec @@ -16,18 +15,34 @@ 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.findSubClasses import dev.inmo.micro_ksp.generator.writeFile import dev.inmo.microutils.kps.sealed.GenerateSealedWorkaround +import java.io.File class Processor( private val codeGenerator: CodeGenerator ) : SymbolProcessor { - private fun KSClassDeclaration.resolveSubclasses(): List { - return (getSealedSubclasses().flatMap { - it.resolveSubclasses() - }.ifEmpty { - sequenceOf(this) - }).toList() + private fun KSClassDeclaration.findSealedConnection(potentialSealedParent: KSClassDeclaration): Boolean { + val targetClassname = potentialSealedParent.qualifiedName ?.asString() + return superTypes.any { + targetClassname == ((it.resolve().declaration as? KSClassDeclaration) ?.qualifiedName ?.asString()) || (it is KSClassDeclaration && it.getSealedSubclasses().any() && it.findSealedConnection(potentialSealedParent)) + } + } + + private fun KSClassDeclaration.resolveSubclasses( + searchIn: Sequence, + allowNonSealed: Boolean + ): Sequence { + return findSubClasses(searchIn).let { + if (allowNonSealed) { + it + } else { + it.filter { + it.findSealedConnection(this) + } + } + } } @OptIn(KspExperimental::class) @@ -35,7 +50,11 @@ class Processor( ksClassDeclaration: KSClassDeclaration, resolver: Resolver ) { - val subClasses = ksClassDeclaration.resolveSubclasses().distinct() + val annotation = ksClassDeclaration.getAnnotationsByType(GenerateSealedWorkaround::class).first() + val subClasses = ksClassDeclaration.resolveSubclasses( + searchIn = resolver.getAllFiles(), + allowNonSealed = annotation.includeNonSealedSubTypes + ).distinct() val subClassesNames = subClasses.filter { when (it.classKind) { ClassKind.ENUM_ENTRY, @@ -51,7 +70,7 @@ class Processor( (it.getAnnotationsByType(GenerateSealedWorkaround.Order::class).firstOrNull()) ?.order ?: 0 }.map { it.toClassName() - } + }.toList() val className = ksClassDeclaration.toClassName() val setType = Set::class.asTypeName().parameterizedBy( ksClassDeclaration.toClassName() diff --git a/ksp/sealed/src/commonMain/kotlin/GenerateSealedWorkaround.kt b/ksp/sealed/src/commonMain/kotlin/GenerateSealedWorkaround.kt index 0dcf29fb50c..c98eed90d8e 100644 --- a/ksp/sealed/src/commonMain/kotlin/GenerateSealedWorkaround.kt +++ b/ksp/sealed/src/commonMain/kotlin/GenerateSealedWorkaround.kt @@ -3,7 +3,8 @@ package dev.inmo.microutils.kps.sealed @Retention(AnnotationRetention.BINARY) @Target(AnnotationTarget.CLASS) annotation class GenerateSealedWorkaround( - val prefix: String = "" + val prefix: String = "", + val includeNonSealedSubTypes: Boolean = false ) { @Retention(AnnotationRetention.BINARY) @Target(AnnotationTarget.CLASS)