Support sealed, generic and abstract polymorphic inheritors in sealed class serialization

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

@@ -90,12 +90,16 @@ fun AbstractSerialGenerator.getSerialTypeInfo(property: SerializableProperty): S
     }
 }
 
-fun AbstractSerialGenerator.immediateSealedSerializableSubclassesFor(
+fun AbstractSerialGenerator.allSealedSerializableSubclassesFor(
     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() }
+    fun recursiveSealed(klass: ClassDescriptor): Collection<ClassDescriptor> {
+        return klass.sealedSubclasses.flatMap { if (it.modality == Modality.SEALED) recursiveSealed(it) else setOf(it) }
+    }
+
+    val serializableSubtypes = recursiveSealed(klass).map { it.toSimpleType() }
     return serializableSubtypes to serializableSubtypes.mapNotNull { findTypeSerializerOrContextUnchecked(module, it) }
 }
 

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

@@ -33,8 +33,8 @@ import org.jetbrains.kotlin.types.*
 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.allSealedSerializableSubclassesFor
 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.*
@@ -466,7 +466,6 @@ interface IrBuilderExtension {
         return serializerInstance(
             generator,
             dispatchReceiverParameter,
-            generator.serializableDescriptor,
             serializer,
             property.module,
             property.type,
@@ -506,37 +505,34 @@ interface IrBuilderExtension {
     fun IrBuilderWithScope.serializerInstance(
         enclosingGenerator: SerializerIrGenerator,
         dispatchReceiverParameter: IrValueParameter,
-        serializableDescriptor: ClassDescriptor,
         serializerClassOriginal: ClassDescriptor?,
         module: ModuleDescriptor,
         kType: KotlinType,
         genericIndex: Int? = null
     ): IrExpression? = serializerInstance(
         enclosingGenerator,
-        serializableDescriptor,
         serializerClassOriginal,
         module,
         kType,
         genericIndex
-    ) {
+    ) { it, _ ->
         val prop = enclosingGenerator.localSerializersFieldsDescriptors[it]
         irGetField(irGet(dispatchReceiverParameter), compilerContext.localSymbolTable.referenceField(prop).owner)
     }
 
     fun IrBuilderWithScope.serializerInstance(
         enclosingGenerator: AbstractSerialGenerator,
-        serializableDescriptor: ClassDescriptor,
         serializerClassOriginal: ClassDescriptor?,
         module: ModuleDescriptor,
         kType: KotlinType,
         genericIndex: Int? = null,
-        genericGetter: (Int) -> IrExpression
+        genericGetter: ((Int, KotlinType) -> IrExpression)? = null
     ): IrExpression? {
         val nullableSerClass =
             compilerContext.externalSymbols.referenceClass(module.getClassFromInternalSerializationPackage(SpecialBuiltins.nullableSerializer))
         if (serializerClassOriginal == null) {
             if (genericIndex == null) return null
-            return genericGetter(genericIndex)
+            return genericGetter?.invoke(genericIndex, kType)
         }
         if (serializerClassOriginal.kind == ClassKind.OBJECT) {
             return irGetObject(serializerClassOriginal)
@@ -544,7 +540,6 @@ interface IrBuilderExtension {
             fun instantiate(serializer: ClassDescriptor?, type: KotlinType): IrExpression? {
                 val expr = serializerInstance(
                     enclosingGenerator,
-                    serializableDescriptor,
                     serializer,
                     module,
                     type,
@@ -569,7 +564,8 @@ interface IrBuilderExtension {
                 sealedSerializerId -> {
                     args = mutableListOf<IrExpression>().apply {
                         add(irString(kType.serialName()))
-                        val (subclasses, subSerializers) = enclosingGenerator.immediateSealedSerializableSubclassesFor(
+                        add(classReference(kType))
+                        val (subclasses, subSerializers) = enclosingGenerator.allSealedSerializableSubclassesFor(
                             kType.toClassDescriptor!!,
                             module
                         )
@@ -583,14 +579,33 @@ interface IrBuilderExtension {
                         add(
                             createArrayOfExpression(
                                 wrapIrTypeIntoKSerializerIrType(module, thisIrType, variance = Variance.OUT_VARIANCE),
-                                subSerializers.mapIndexed { i, ser -> instantiate(ser, subclasses[i])!! }
+                                subSerializers.mapIndexed { i, serializer ->
+                                    val type = subclasses[i]
+                                    val expr = serializerInstance(
+                                        enclosingGenerator,
+                                        serializer,
+                                        module,
+                                        type,
+                                        type.genericIndex
+                                    ) { _, genericType ->
+                                        serializerInstance(
+                                            enclosingGenerator,
+                                            module.getClassFromSerializationPackage(
+                                                SpecialBuiltins.polymorphicSerializer
+                                            ),
+                                            module,
+                                            genericType
+                                        )!!
+                                    }!!
+                                    wrapWithNullableSerializerIfNeeded(module, type, expr, nullableSerClass)
+                                }
                             )
                         )
                     }
                     typeArgs = listOf(thisIrType)
                 }
                 enumSerializerId -> {
-                    serializerClass = serializableDescriptor.getClassFromInternalSerializationPackage("CommonEnumSerializer")
+                    serializerClass = module.getClassFromInternalSerializationPackage("CommonEnumSerializer")
                     args = kType.toClassDescriptor!!.let { enumDesc ->
                         listOf(
                             irString(enumDesc.serialName()),

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

@@ -11,7 +11,6 @@ import org.jetbrains.kotlin.ir.builders.irReturn
 import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.impl.IrConstructorCallImpl
-import org.jetbrains.kotlin.ir.types.defaultType
 import org.jetbrains.kotlin.ir.util.SymbolTable
 import org.jetbrains.kotlin.ir.util.TypeTranslator
 import org.jetbrains.kotlin.ir.util.defaultType
@@ -23,8 +22,6 @@ import org.jetbrains.kotlin.resolve.descriptorUtil.module
 import org.jetbrains.kotlinx.serialization.compiler.backend.common.SerializableCompanionCodegen
 import org.jetbrains.kotlinx.serialization.compiler.backend.common.findTypeSerializer
 import org.jetbrains.kotlinx.serialization.compiler.resolve.*
-import org.jetbrains.kotlinx.serialization.compiler.resolve.getSerializableClassDescriptorByCompanion
-import org.jetbrains.kotlinx.serialization.compiler.resolve.shouldHaveGeneratedMethodsInCompanion
 
 class SerializableCompanionIrGenerator(
     val irClass: IrClass,
@@ -94,10 +91,9 @@ class SerializableCompanionIrGenerator(
             val args: List<IrExpression> = getter.valueParameters.map { irGet(it) }
             val expr = serializerInstance(
                 this@SerializableCompanionIrGenerator,
-                serializableDescriptor, serializer,
+                serializer, serializableDescriptor.module,
-                serializableDescriptor.module,
                 serializableDescriptor.defaultType
-            ) { args[it] }
+            ) { it, _ -> args[it] }
             patchSerializableClassWithMarkerAnnotation(serializer)
             +irReturn(requireNotNull(expr))
         }

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

@@ -248,7 +248,6 @@ open class SerializerIrGenerator(val irClass: IrClass, final override val compil
             val innerSerial = serializerInstance(
                 this@SerializerIrGenerator,
                 saveFunc.dispatchReceiverParameter!!,
-                serializableDescriptor,
                 sti.serializer,
                 property.module,
                 property.type,
@@ -381,7 +380,6 @@ open class SerializerIrGenerator(val irClass: IrClass, final override val compil
                                 val innerSerial = serializerInstance(
                                     this@SerializerIrGenerator,
                                     loadFunc.dispatchReceiverParameter!!,
-                                    serializableDescriptor,
                                     sti.serializer,
                                     property.module,
                                     property.type,

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

@@ -30,6 +30,7 @@ import org.jetbrains.kotlin.resolve.descriptorUtil.classId
 import org.jetbrains.kotlin.resolve.descriptorUtil.getSuperClassNotAny
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.typeUtil.isTypeParameter
+import org.jetbrains.kotlin.types.typeUtil.representativeUpperBound
 import org.jetbrains.kotlinx.serialization.compiler.backend.common.*
 import org.jetbrains.kotlinx.serialization.compiler.backend.jvm.*
 import org.jetbrains.kotlinx.serialization.compiler.resolve.*
@@ -134,13 +135,18 @@ internal fun AbstractSerialGenerator.serializerInstance(
     module: ModuleDescriptor,
     kType: KotlinType,
     genericIndex: Int? = null,
-    genericGetter: (Int) -> JsExpression = { JsNameRef(context.scope().declareName("${SerialEntityNames.typeArgPrefix}$it"), JsThisRef()) }
+    genericGetter: (Int, KotlinType) -> JsExpression = { it, _ ->
+        JsNameRef(
+            context.scope().declareName("${SerialEntityNames.typeArgPrefix}$it"),
+            JsThisRef()
+        )
+    }
 ): JsExpression? {
     val nullableSerClass =
         context.translateQualifiedReference(module.getClassFromInternalSerializationPackage(SpecialBuiltins.nullableSerializer))
     if (serializerClass == null) {
         if (genericIndex == null) return null
-        return genericGetter(genericIndex)
+        return genericGetter(genericIndex, kType)
     }
     if (serializerClass.kind == ClassKind.OBJECT) {
         return context.serializerObjectGetter(serializerClass)
@@ -163,7 +169,8 @@ internal fun AbstractSerialGenerator.serializerInstance(
             )
             serializerClass.classId == sealedSerializerId -> mutableListOf<JsExpression>().apply {
                 add(JsStringLiteral(kType.serialName()))
-                val (subclasses, subSerializers) = immediateSealedSerializableSubclassesFor(
+                add(ExpressionVisitor.getObjectKClass(context, kType.toClassDescriptor!!))
+                val (subclasses, subSerializers) = allSealedSerializableSubclassesFor(
                     kType.toClassDescriptor!!,
                     module
                 )
@@ -173,8 +180,17 @@ internal fun AbstractSerialGenerator.serializerInstance(
                         it.toClassDescriptor!!
                     )
                 }))
-                add(JsArrayLiteral(subSerializers.mapIndexed { i, ser ->
+                add(JsArrayLiteral(subSerializers.mapIndexed { i, serializer ->
-                    instantiate(ser, subclasses[i])!!
+                    val type = subclasses[i]
+                    val expr = serializerInstance(context, serializer, module, type, type.genericIndex) { _, genericType ->
+                        serializerInstance(
+                            context,
+                            module.getClassFromSerializationPackage(SpecialBuiltins.polymorphicSerializer),
+                            module,
+                            (genericType.constructor.declarationDescriptor as TypeParameterDescriptor).representativeUpperBound
+                        )!!
+                    }!!
+                    if (type.isMarkedNullable) JsNew(nullableSerClass, listOf(expr)) else expr
                 }))
             }
             else -> kType.arguments.map {

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

@@ -52,7 +52,7 @@ class SerializableCompanionJsTranslator(
                         serializer,
                         serializableDescriptor.module,
                         serializableDescriptor.defaultType,
-                        genericGetter = {
+                        genericGetter = { it, _ ->
                             args[it]
                         })
                 )

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

@@ -29,6 +29,7 @@ import org.jetbrains.kotlin.resolve.jvm.diagnostics.OtherOrigin
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.typeUtil.isTypeParameter
+import org.jetbrains.kotlin.types.typeUtil.representativeUpperBound
 import org.jetbrains.kotlinx.serialization.compiler.backend.common.*
 import org.jetbrains.kotlinx.serialization.compiler.resolve.*
 import org.jetbrains.kotlinx.serialization.compiler.resolve.SerialEntityNames.DECODER_CLASS
@@ -177,7 +178,7 @@ internal fun InstructionAdapter.stackValueSerializerInstanceFromClass(
         serializer,
         this,
         sti.property.genericIndex
-    ) { idx ->
+    ) { idx, _ ->
         load(varIndexStart + idx, kSerializerType)
     }
 }
@@ -201,7 +202,7 @@ internal fun InstructionAdapter.stackValueSerializerInstanceFromSerializerWithou
         serializer,
         this,
         property.genericIndex
-    ) { idx ->
+    ) { idx, _ ->
         load(0, kSerializerType)
         getfield(codegen.typeMapper.mapClass(codegen.descriptor).internalName, "$typeArgPrefix$idx", kSerializerType.descriptor)
     }.also { if (it && property.type.isMarkedNullable) wrapStackValueIntoNullableSerializer() }
@@ -211,7 +212,7 @@ internal fun InstructionAdapter.stackValueSerializerInstanceFromSerializer(codeg
     return serializerCodegen.stackValueSerializerInstance(
         codegen, sti.property.module, sti.property.type,
         sti.serializer, this, sti.property.genericIndex
-    ) { idx ->
+    ) { idx, _ ->
         load(0, kSerializerType)
         getfield(codegen.typeMapper.mapClass(codegen.descriptor).internalName, "$typeArgPrefix$idx", kSerializerType.descriptor)
     }
@@ -220,10 +221,13 @@ internal fun InstructionAdapter.stackValueSerializerInstanceFromSerializer(codeg
 // returns false is cannot not use serializer
 // use iv == null to check only (do not emit serializer onto stack)
 internal fun AbstractSerialGenerator.stackValueSerializerInstance(codegen: ClassBodyCodegen, module: ModuleDescriptor, kType: KotlinType, maybeSerializer: ClassDescriptor?,
-                                          iv: InstructionAdapter?, genericIndex: Int? = null, genericSerializerFieldGetter: (InstructionAdapter.(Int) -> Unit)? = null): Boolean {
+                                                                  iv: InstructionAdapter?,
+                                                                  genericIndex: Int? = null,
+                                                                  genericSerializerFieldGetter: (InstructionAdapter.(Int, KotlinType) -> Unit)? = null
+): Boolean {
     if (maybeSerializer == null && genericIndex != null) {
         // get field from serializer object
-        iv?.run { genericSerializerFieldGetter?.invoke(this, genericIndex) }
+        iv?.run { genericSerializerFieldGetter?.invoke(this, genericIndex, kType) }
         return true
     }
     val serializer = maybeSerializer ?: return false
@@ -302,16 +306,41 @@ internal fun AbstractSerialGenerator.stackValueSerializerInstance(codegen: Class
             sealedSerializerId -> {
                 aconst(serialName)
                 signature.append("Ljava/lang/String;")
-                val (subClasses, subSerializers) = immediateSealedSerializableSubclassesFor(kType.toClassDescriptor!!, module)
+                aconst(codegen.typeMapper.mapType(kType, null, TypeMappingMode.GENERIC_ARGUMENT))
+                AsmUtil.wrapJavaClassIntoKClass(this)
+                signature.append(AsmTypes.K_CLASS_TYPE.descriptor)
+                val (subClasses, subSerializers) = allSealedSerializableSubclassesFor(kType.toClassDescriptor!!, module)
                 // KClasses vararg
-                fillArray(AsmTypes.K_CLASS_TYPE, subClasses) { i, type ->
+                fillArray(AsmTypes.K_CLASS_TYPE, subClasses) { _, 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)
+                    val (argType, argSerializer) = subClasses[i] to serializer
+                    assert(
+                        stackValueSerializerInstance(
+                            codegen,
+                            module,
+                            argType,
+                            argSerializer,
+                            this,
+                            argType.genericIndex
+                        ) { _, genericType ->
+                            // if we encountered generic type parameter in one of subclasses of sealed class, use polymorphism from upper bound
+                            assert(
+                                stackValueSerializerInstance(
+                                    codegen,
+                                    module,
+                                    (genericType.constructor.declarationDescriptor as TypeParameterDescriptor).representativeUpperBound,
+                                    module.getClassFromSerializationPackage(SpecialBuiltins.polymorphicSerializer),
+                                    this
+                                )
+                            )
+                        }
+                    )
+                    if (argType.isMarkedNullable) wrapStackValueIntoNullableSerializer()
                 }
                 signature.append(kSerializerArrayType.descriptor)
             }

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

@@ -104,8 +104,8 @@ class SerializableCodegenImpl(
             superTypeArguments.forEach {
                 val genericIdx = serializableDescriptor.defaultType.arguments.indexOf(it).let { if (it == -1) null else it }
                 val serial = findTypeSerializerOrContext(serializableDescriptor.module, it.type)
-                stackValueSerializerInstance(classCodegen, serializableDescriptor.module, it.type, serial, this, genericIdx) {
+                stackValueSerializerInstance(classCodegen, serializableDescriptor.module, it.type, serial, this, genericIdx) { i, _ ->
-                    load(offsetI + it, kSerializerType)
+                    load(offsetI + i, kSerializerType)
                 }
             }
             val superSignature =

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

@@ -51,7 +51,7 @@ class SerializableCompanionCodegenImpl(private val classCodegen: ImplementationB
                 serial,
                 this,
                 null
-            ) {
+            ) { it, _ ->
                 load(it + 1, kSerializerType)
             }
             areturn(kSerializerType)