Sealed classes serialization

plugins/kotlin-serialization/kotlin-serialization-compiler/src/org/jetbrains/kotlinx/serialization/compiler/backend/common/SerializableCodegen.kt

@@ -36,7 +36,7 @@ abstract class SerializableCodegen(
 
     private fun generateSyntheticInternalConstructor() {
         val serializerDescriptor = serializableDescriptor.classSerializer ?: return
-        if (isAbstractSerializableClass(serializableDescriptor) || SerializerCodegen.getSyntheticLoadMember(serializerDescriptor) != null) {
+        if (serializableDescriptor.isAbstractSerializableClass() || serializableDescriptor.isSealedSerializableClass() || SerializerCodegen.getSyntheticLoadMember(serializerDescriptor) != null) {
             val constrDesc = serializableDescriptor.secondaryConstructors.find(ClassConstructorDescriptor::isSerializationCtor) ?: return
             generateInternalConstructor(constrDesc)
         }
@@ -44,7 +44,7 @@ abstract class SerializableCodegen(
 
     private fun generateSyntheticMethods() {
         val serializerDescriptor = serializableDescriptor.classSerializer ?: return
-        if (isAbstractSerializableClass(serializableDescriptor) || SerializerCodegen.getSyntheticSaveMember(serializerDescriptor) != null) {
+        if (serializableDescriptor.isAbstractSerializableClass() || serializableDescriptor.isSealedSerializableClass() || SerializerCodegen.getSyntheticSaveMember(serializerDescriptor) != null) {
             val func = KSerializerDescriptorResolver.createWriteSelfFunctionDescriptor(serializableDescriptor)
             generateWriteSelfMethod(func)
         }

plugins/kotlin-serialization/kotlin-serialization-compiler/src/org/jetbrains/kotlinx/serialization/compiler/backend/common/TypeUtil.kt

@@ -90,6 +90,15 @@ fun AbstractSerialGenerator.getSerialTypeInfo(property: SerializableProperty): S
     }
 }
 
+fun AbstractSerialGenerator.immediateSealedSerializableSubclassesFor(
+    klass: ClassDescriptor,
+    module: ModuleDescriptor
+): Pair<List<KotlinType>, List<ClassDescriptor>> {
+    assert(klass.kind == ClassKind.CLASS && klass.modality == Modality.SEALED)
+    val serializableSubtypes = klass.sealedSubclasses.map { it.toSimpleType() }
+    return serializableSubtypes to serializableSubtypes.mapNotNull { findTypeSerializerOrContextUnchecked(module, it) }
+}
+
 fun KotlinType.serialName(): String {
     val serializableDescriptor = this.toClassDescriptor!!
     return serializableDescriptor.serialName()

plugins/kotlin-serialization/kotlin-serialization-compiler/src/org/jetbrains/kotlinx/serialization/compiler/backend/ir/GeneratorHelpers.kt

@@ -27,19 +27,14 @@ import org.jetbrains.kotlin.ir.types.toKotlinType
 import org.jetbrains.kotlin.ir.types.typeWith
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.js.resolve.diagnostics.findPsi
-import org.jetbrains.kotlin.name.ClassId
-import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.descriptorUtil.classId
-import org.jetbrains.kotlin.resolve.descriptorUtil.module
+import org.jetbrains.kotlin.types.*
-import org.jetbrains.kotlin.types.KotlinType
-import org.jetbrains.kotlin.types.KotlinTypeFactory
-import org.jetbrains.kotlin.types.TypeProjectionImpl
-import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.types.typeUtil.isTypeParameter
 import org.jetbrains.kotlin.types.typeUtil.makeNotNullable
 import org.jetbrains.kotlinx.serialization.compiler.backend.common.AbstractSerialGenerator
 import org.jetbrains.kotlinx.serialization.compiler.backend.common.findTypeSerializerOrContext
+import org.jetbrains.kotlinx.serialization.compiler.backend.common.immediateSealedSerializableSubclassesFor
 import org.jetbrains.kotlinx.serialization.compiler.backend.common.serialName
 import org.jetbrains.kotlinx.serialization.compiler.backend.jvm.*
 import org.jetbrains.kotlinx.serialization.compiler.resolve.*
@@ -399,11 +394,15 @@ interface IrBuilderExtension {
         }
     }
 
+    fun kClassTypeFor(projection: TypeProjection): SimpleType {
+        val kClass = compilerContext.builtIns.kClass
+        return KotlinTypeFactory.simpleNotNullType(Annotations.EMPTY, kClass, listOf(projection))
+    }
+
     fun createClassReference(classType: KotlinType, startOffset: Int, endOffset: Int): IrClassReference {
         val clazz = classType.toClassDescriptor!!
-        val kClass = clazz.module.findClassAcrossModuleDependencies(ClassId(FqName("kotlin.reflect"), Name.identifier("KClass")))!!
         val returnType =
-            KotlinTypeFactory.simpleNotNullType(Annotations.EMPTY, kClass, listOf(TypeProjectionImpl(Variance.INVARIANT, classType)))
+            kClassTypeFor(TypeProjectionImpl(Variance.INVARIANT, classType))
         return IrClassReferenceImpl(
             startOffset,
             endOffset,
@@ -494,12 +493,12 @@ interface IrBuilderExtension {
     }
 
 
-    fun wrapIrTypeIntoKSerializerIrType(module: ModuleDescriptor, type: IrType): IrType {
+    fun wrapIrTypeIntoKSerializerIrType(module: ModuleDescriptor, type: IrType, variance: Variance = Variance.INVARIANT): IrType {
         val kSerClass =
             compilerContext.externalSymbols.referenceClass(module.getClassFromSerializationPackage(SerialEntityNames.KSERIALIZER_CLASS))
         return IrSimpleTypeImpl(
             kSerClass, hasQuestionMark = false, arguments = listOf(
-                makeTypeProjection(type, Variance.INVARIANT)
+                makeTypeProjection(type, variance)
             ), annotations = emptyList()
         )
     }
@@ -542,17 +541,53 @@ interface IrBuilderExtension {
         if (serializerClassOriginal.kind == ClassKind.OBJECT) {
             return irGetObject(serializerClassOriginal)
         } else {
+            fun instantiate(serializer: ClassDescriptor?, type: KotlinType): IrExpression? {
+                val expr = serializerInstance(
+                    enclosingGenerator,
+                    serializableDescriptor,
+                    serializer,
+                    module,
+                    type,
+                    type.genericIndex,
+                    genericGetter
+                ) ?: return null
+                return wrapWithNullableSerializerIfNeeded(module, type, expr, nullableSerClass)
+            }
             var serializerClass = serializerClassOriginal
             var args: List<IrExpression>
             var typeArgs: List<IrType?>
+            val thisIrType = kType.toIrType()
             when (serializerClassOriginal.classId) {
                 contextSerializerId, polymorphicSerializerId -> {
                     args = listOf(classReference(kType))
-                    typeArgs = listOf(kType.toIrType())
+                    typeArgs = listOf(thisIrType)
                 }
                 objectSerializerId -> {
                     args = listOf(irString(kType.serialName()), irGetObject(kType.toClassDescriptor!!))
-                    typeArgs = listOf(kType.toIrType())
+                    typeArgs = listOf(thisIrType)
+                }
+                sealedSerializerId -> {
+                    args = mutableListOf<IrExpression>().apply {
+                        add(irString(kType.serialName()))
+                        val (subclasses, subSerializers) = enclosingGenerator.immediateSealedSerializableSubclassesFor(
+                            kType.toClassDescriptor!!,
+                            module
+                        )
+                        val projectedOutCurrentKClass = kClassTypeFor(TypeProjectionImpl(Variance.OUT_VARIANCE, kType))
+                        add(
+                            createArrayOfExpression(
+                                projectedOutCurrentKClass.toIrType(),
+                                subclasses.map { classReference(it) }
+                            )
+                        )
+                        add(
+                            createArrayOfExpression(
+                                wrapIrTypeIntoKSerializerIrType(module, thisIrType, variance = Variance.OUT_VARIANCE),
+                                subSerializers.mapIndexed { i, ser -> instantiate(ser, subclasses[i])!! }
+                            )
+                        )
+                    }
+                    typeArgs = listOf(thisIrType)
                 }
                 enumSerializerId -> {
                     serializerClass = serializableDescriptor.getClassFromInternalSerializationPackage("CommonEnumSerializer")
@@ -566,7 +601,7 @@ interface IrBuilderExtension {
                             )
                         )
                     }
-                    typeArgs = listOf(kType.toIrType())
+                    typeArgs = listOf(thisIrType)
                 }
                 else -> {
                     args = kType.arguments.map {
@@ -575,17 +610,7 @@ interface IrBuilderExtension {
                             it.type,
                             sourceElement = serializerClassOriginal.findPsi()
                         )
-                        val expr = serializerInstance(
+                        instantiate(argSer, it.type) ?: return null
-                            enclosingGenerator,
-                            serializableDescriptor,
-                            argSer,
-                            module,
-                            it.type,
-                            it.type.genericIndex,
-                            genericGetter
-                        )
-                            ?: return null
-                        wrapWithNullableSerializerIfNeeded(module, it.type, expr, nullableSerClass)
                     }
                     typeArgs = kType.arguments.map { it.type.toIrType() }
                 }
@@ -606,8 +631,7 @@ interface IrBuilderExtension {
             } else {
                 compilerContext.externalSymbols.referenceConstructor(serializerClass.unsubstitutedPrimaryConstructor!!)
             }
-            val propertyTypeHint = kType.toIrType()
+            val returnType = wrapIrTypeIntoKSerializerIrType(module, thisIrType)
-            val returnType = wrapIrTypeIntoKSerializerIrType(module, propertyTypeHint)
             return irInvoke(null, ctor, typeArguments = typeArgs, valueArguments = args, returnTypeHint = returnType)
         }
     }

plugins/kotlin-serialization/kotlin-serialization-compiler/src/org/jetbrains/kotlinx/serialization/compiler/backend/js/JsCodegenUtil.kt

@@ -145,6 +145,10 @@ internal fun AbstractSerialGenerator.serializerInstance(
     if (serializerClass.kind == ClassKind.OBJECT) {
         return context.serializerObjectGetter(serializerClass)
     } else {
+        fun instantiate(serializer: ClassDescriptor?, type: KotlinType): JsExpression? {
+            val expr = serializerInstance(context, serializer, module, type, type.genericIndex, genericGetter) ?: return null
+            return if (type.isMarkedNullable) JsNew(nullableSerClass, listOf(expr)) else expr
+        }
         var args = when {
             serializerClass.classId == contextSerializerId || serializerClass.classId == polymorphicSerializerId -> listOf(
                 ExpressionVisitor.getObjectKClass(context, kType.toClassDescriptor!!)
@@ -157,10 +161,25 @@ internal fun AbstractSerialGenerator.serializerInstance(
                 JsStringLiteral(kType.serialName()),
                 context.serializerObjectGetter(kType.toClassDescriptor!!)
             )
+            serializerClass.classId == sealedSerializerId -> mutableListOf<JsExpression>().apply {
+                add(JsStringLiteral(kType.serialName()))
+                val (subclasses, subSerializers) = immediateSealedSerializableSubclassesFor(
+                    kType.toClassDescriptor!!,
+                    module
+                )
+                add(JsArrayLiteral(subclasses.map {
+                    ExpressionVisitor.getObjectKClass(
+                        context,
+                        it.toClassDescriptor!!
+                    )
+                }))
+                add(JsArrayLiteral(subSerializers.mapIndexed { i, ser ->
+                    instantiate(ser, subclasses[i])!!
+                }))
+            }
             else -> kType.arguments.map {
                 val argSer = findTypeSerializerOrContext(module, it.type, sourceElement = serializerClass.findPsi())
-                val expr = serializerInstance(context, argSer, module, it.type, it.type.genericIndex, genericGetter) ?: return null
+                instantiate(argSer, it.type) ?: return null
-                if (it.type.isMarkedNullable) JsNew(nullableSerClass, listOf(expr)) else expr
             }
         }
         if (serializerClass.classId == referenceArraySerializerId)

plugins/kotlin-serialization/kotlin-serialization-compiler/src/org/jetbrains/kotlinx/serialization/compiler/backend/jvm/JVMCodegenUtil.kt

@@ -159,6 +159,7 @@ internal val enumSerializerId = ClassId(internalPackageFqName, Name.identifier(S
 internal val polymorphicSerializerId = ClassId(packageFqName, Name.identifier(SpecialBuiltins.polymorphicSerializer))
 internal val referenceArraySerializerId = ClassId(internalPackageFqName, Name.identifier(SpecialBuiltins.referenceArraySerializer))
 internal val objectSerializerId = ClassId(internalPackageFqName, Name.identifier(SpecialBuiltins.objectSerializer))
+internal val sealedSerializerId = ClassId(packageFqName, Name.identifier(SpecialBuiltins.sealedSerializer))
 internal val contextSerializerId = ClassId(packageFqName, Name.identifier(SpecialBuiltins.contextSerializer))
 
 
@@ -263,7 +264,7 @@ internal fun AbstractSerialGenerator.stackValueSerializerInstance(codegen: Class
         // instantiate all arg serializers on stack
         val signature = StringBuilder("(")
 
-        fun instantiate(typeArgument: Pair<KotlinType, ClassDescriptor?>) {
+        fun instantiate(typeArgument: Pair<KotlinType, ClassDescriptor?>, writeSignature: Boolean = true) {
             val (argType, argSerializer) = typeArgument
             assert(
                 stackValueSerializerInstance(
@@ -278,9 +279,10 @@ internal fun AbstractSerialGenerator.stackValueSerializerInstance(codegen: Class
             )
             // wrap into nullable serializer if argType is nullable
             if (argType.isMarkedNullable) wrapStackValueIntoNullableSerializer()
-            signature.append(kSerializerType.descriptor)
+            if (writeSignature) signature.append(kSerializerType.descriptor)
         }
 
+        val serialName = kType.serialName()
         when (serializer.classId) {
             enumSerializerId, contextSerializerId, polymorphicSerializerId -> {
                 // a special way to instantiate enum -- need a enum KClass reference
@@ -297,15 +299,30 @@ internal fun AbstractSerialGenerator.stackValueSerializerInstance(codegen: Class
                 // Reference array serializer still needs serializer for its argument type
                 instantiate(argSerializers[0])
             }
+            sealedSerializerId -> {
+                aconst(serialName)
+                signature.append("Ljava/lang/String;")
+                val (subClasses, subSerializers) = immediateSealedSerializableSubclassesFor(kType.toClassDescriptor!!, module)
+                // KClasses vararg
+                fillArray(AsmTypes.K_CLASS_TYPE, subClasses) { i, type ->
+                    aconst(codegen.typeMapper.mapType(type, null, TypeMappingMode.GENERIC_ARGUMENT))
+                    AsmUtil.wrapJavaClassIntoKClass(this)
+                }
+                signature.append(AsmTypes.K_CLASS_ARRAY_TYPE.descriptor)
+                // Serializers vararg
+                fillArray(kSerializerType, subSerializers) { i, serializer ->
+                    instantiate(subClasses[i] to serializer, writeSignature = false)
+                }
+                signature.append(kSerializerArrayType.descriptor)
+            }
             objectSerializerId -> {
-                val serialName = kType.serialName()
                 aconst(serialName)
                 signature.append("Ljava/lang/String;")
                 StackValue.singleton(kType.toClassDescriptor!!, codegen.typeMapper).put(Type.getType("Ljava/lang/Object;"), iv)
                 signature.append("Ljava/lang/Object;")
             }
             // all serializers get arguments with serializers of their generic types
-            else -> argSerializers.forEach(::instantiate)
+            else -> argSerializers.forEach { instantiate(it) }
         }
         signature.append(")V")
         // invoke constructor
@@ -320,6 +337,17 @@ fun InstructionAdapter.wrapStackValueIntoNullableSerializer() =
         "(" + kSerializerType.descriptor + ")" + kSerializerType.descriptor, false
     )
 
+fun <T> InstructionAdapter.fillArray(type: Type, args: List<T>, onEach: (Int, T) -> Unit) {
+    iconst(args.size)
+    newarray(type)
+    args.forEachIndexed { i, serializer ->
+        dup()
+        iconst(i)
+        onEach(i, serializer)
+        astore(type)
+    }
+}
+
 //
 // ======= Serializers Resolving =======
 //

plugins/kotlin-serialization/kotlin-serialization-compiler/src/org/jetbrains/kotlinx/serialization/compiler/resolve/KSerializationUtil.kt

@@ -28,9 +28,8 @@ import org.jetbrains.kotlin.resolve.descriptorUtil.module
 import org.jetbrains.kotlin.resolve.lazy.descriptors.LazyAnnotationDescriptor
 import org.jetbrains.kotlin.resolve.scopes.getDescriptorsFiltered
 import org.jetbrains.kotlin.types.*
-import org.jetbrains.kotlin.types.typeUtil.*
+import org.jetbrains.kotlin.types.typeUtil.isSubtypeOf
-import org.jetbrains.kotlinx.serialization.compiler.backend.jvm.*
+import org.jetbrains.kotlin.types.typeUtil.representativeUpperBound
-import org.jetbrains.kotlinx.serialization.compiler.resolve.SerialEntityNames.KSERIALIZER_CLASS
 import org.jetbrains.kotlinx.serialization.compiler.resolve.SerializationAnnotations.serialInfoFqName
 
 internal fun isAllowedToHaveAutoGeneratedSerializerMethods(
@@ -128,7 +127,8 @@ internal fun ClassDescriptor.isSerializableEnum(): Boolean {
     return false
 }
 
-internal val ClassDescriptor.shouldHaveGeneratedSerializer: Boolean get() = isInternalSerializable || isSerializableEnum()
+internal val ClassDescriptor.shouldHaveGeneratedSerializer: Boolean
+    get() = (isInternalSerializable && (modality == Modality.FINAL || modality == Modality.OPEN)) || isSerializableEnum()
 
 internal fun ClassDescriptor.enumEntries(): List<ClassDescriptor> {
     check(this.kind == ClassKind.ENUM_CLASS)
@@ -158,14 +158,21 @@ internal fun Annotated.findSerializableAnnotationDeclaration(): KtAnnotationEntr
 // For abstract classes marked with @Serializable,
 // methods are generated anyway although they shouldn't have
 // generated $serializer and use Polymorphic one.
-internal fun isAbstractSerializableClass(serializableDescriptor: ClassDescriptor): Boolean =
+internal fun ClassDescriptor.isAbstractSerializableClass(): Boolean =
-    serializableDescriptor.isInternalSerializable && serializableDescriptor.modality == Modality.ABSTRACT
+    isInternalSerializable && modality == Modality.ABSTRACT
 
-internal fun ClassDescriptor.polymorphicSerializerIfApplicableAutomatically(): ClassDescriptor? = if (
+internal fun ClassDescriptor.isSealedSerializableClass(): Boolean =
-    isAbstractSerializableClass(this)
+    isInternalSerializable && modality == Modality.SEALED
-    || kind == ClassKind.INTERFACE
+
-) module.getClassFromSerializationPackage(SpecialBuiltins.polymorphicSerializer)
+internal fun ClassDescriptor.polymorphicSerializerIfApplicableAutomatically(): ClassDescriptor? {
-else null
+    val serializer = when {
+        this.isAbstractSerializableClass()
+                || kind == ClassKind.INTERFACE -> SpecialBuiltins.polymorphicSerializer
+        this.isSealedSerializableClass() -> SpecialBuiltins.sealedSerializer
+        else -> null
+    }
+    return serializer?.let { module.getClassFromSerializationPackage(it) }
+}
 
 // serializer that was declared for this type
 internal val ClassDescriptor?.classSerializer: ClassDescriptor?

plugins/kotlin-serialization/kotlin-serialization-compiler/src/org/jetbrains/kotlinx/serialization/compiler/resolve/KSerializerDescriptorResolver.kt

@@ -291,7 +291,7 @@ object KSerializerDescriptorResolver {
         )
 
         val markerDesc = classDescriptor.getKSerializerConstructorMarker()
-        val markerType = markerDesc.toSimpleType()
+        val markerType = markerDesc.toSimpleType(nullable = true)
 
         val serializableProperties = bindingContext.serializablePropertiesFor(classDescriptor).serializableProperties
         val parameterDescsAsProps = serializableProperties.map { it.descriptor }

plugins/kotlin-serialization/kotlin-serialization-compiler/src/org/jetbrains/kotlinx/serialization/compiler/resolve/NamingConventions.kt

@@ -80,6 +80,7 @@ object SpecialBuiltins {
     const val objectSerializer = "ObjectSerializer"
     const val enumSerializer = "EnumSerializer"
     const val polymorphicSerializer = "PolymorphicSerializer"
+    const val sealedSerializer = "SealedClassSerializer"
     const val contextSerializer = "ContextSerializer"
     const val nullableSerializer = "NullableSerializer"
 }

plugins/kotlin-serialization/kotlin-serialization-compiler/src/org/jetbrains/kotlinx/serialization/compiler/resolve/SearchUtils.kt

@@ -91,7 +91,7 @@ internal fun ClassDescriptor.getClassFromSerializationPackage(classSimpleName: S
 internal fun ClassDescriptor.getClassFromInternalSerializationPackage(classSimpleName: String) =
     module.getClassFromInternalSerializationPackage(classSimpleName)
 
-fun ClassDescriptor.toSimpleType(nullable: Boolean = true) =
+fun ClassDescriptor.toSimpleType(nullable: Boolean = false) =
     KotlinTypeFactory.simpleType(Annotations.EMPTY, this.typeConstructor, emptyList(), nullable)
 
 internal fun Annotated.annotationsWithArguments(): List<Triple<ClassDescriptor, List<ValueArgument>, List<ValueParameterDescriptor>>> =