improve sealed workaround generation

This commit is contained in:
InsanusMokrassar 2024-07-25 02:27:33 +06:00
parent d164813bb4
commit 698ed6718d
3 changed files with 60 additions and 10 deletions

View File

@ -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<KSClassDeclaration> {
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<KSAnnotated>): Sequence<KSClassDeclaration> {
return files.flatMap {
findSubClasses(it)
}
}

View File

@ -6,7 +6,6 @@ import com.google.devtools.ksp.processing.CodeGenerator
import com.google.devtools.ksp.processing.Resolver import com.google.devtools.ksp.processing.Resolver
import com.google.devtools.ksp.processing.SymbolProcessor import com.google.devtools.ksp.processing.SymbolProcessor
import com.google.devtools.ksp.symbol.* import com.google.devtools.ksp.symbol.*
import com.squareup.kotlinpoet.AnnotationSpec
import com.squareup.kotlinpoet.ClassName import com.squareup.kotlinpoet.ClassName
import com.squareup.kotlinpoet.CodeBlock import com.squareup.kotlinpoet.CodeBlock
import com.squareup.kotlinpoet.FileSpec 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.PropertySpec
import com.squareup.kotlinpoet.asTypeName import com.squareup.kotlinpoet.asTypeName
import com.squareup.kotlinpoet.ksp.toClassName import com.squareup.kotlinpoet.ksp.toClassName
import dev.inmo.micro_ksp.generator.findSubClasses
import dev.inmo.micro_ksp.generator.writeFile import dev.inmo.micro_ksp.generator.writeFile
import dev.inmo.microutils.kps.sealed.GenerateSealedWorkaround import dev.inmo.microutils.kps.sealed.GenerateSealedWorkaround
import java.io.File
class Processor( class Processor(
private val codeGenerator: CodeGenerator private val codeGenerator: CodeGenerator
) : SymbolProcessor { ) : SymbolProcessor {
private fun KSClassDeclaration.resolveSubclasses(): List<KSClassDeclaration> { private fun KSClassDeclaration.findSealedConnection(potentialSealedParent: KSClassDeclaration): Boolean {
return (getSealedSubclasses().flatMap { val targetClassname = potentialSealedParent.qualifiedName ?.asString()
it.resolveSubclasses() return superTypes.any {
}.ifEmpty { targetClassname == ((it.resolve().declaration as? KSClassDeclaration) ?.qualifiedName ?.asString()) || (it is KSClassDeclaration && it.getSealedSubclasses().any() && it.findSealedConnection(potentialSealedParent))
sequenceOf(this) }
}).toList() }
private fun KSClassDeclaration.resolveSubclasses(
searchIn: Sequence<KSAnnotated>,
allowNonSealed: Boolean
): Sequence<KSClassDeclaration> {
return findSubClasses(searchIn).let {
if (allowNonSealed) {
it
} else {
it.filter {
it.findSealedConnection(this)
}
}
}
} }
@OptIn(KspExperimental::class) @OptIn(KspExperimental::class)
@ -35,7 +50,11 @@ class Processor(
ksClassDeclaration: KSClassDeclaration, ksClassDeclaration: KSClassDeclaration,
resolver: Resolver 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 { val subClassesNames = subClasses.filter {
when (it.classKind) { when (it.classKind) {
ClassKind.ENUM_ENTRY, ClassKind.ENUM_ENTRY,
@ -51,7 +70,7 @@ class Processor(
(it.getAnnotationsByType(GenerateSealedWorkaround.Order::class).firstOrNull()) ?.order ?: 0 (it.getAnnotationsByType(GenerateSealedWorkaround.Order::class).firstOrNull()) ?.order ?: 0
}.map { }.map {
it.toClassName() it.toClassName()
} }.toList()
val className = ksClassDeclaration.toClassName() val className = ksClassDeclaration.toClassName()
val setType = Set::class.asTypeName().parameterizedBy( val setType = Set::class.asTypeName().parameterizedBy(
ksClassDeclaration.toClassName() ksClassDeclaration.toClassName()

View File

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