Support annotations in constructor signatures of special serializers

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

@@ -810,6 +810,27 @@ interface IrBuilderExtension {
             .map { it.name.toString() }
     }
 
+    fun IrBuilderWithScope.copyAnnotationsFrom(annotations: List<IrConstructorCall>): List<IrExpression> =
+        annotations.mapNotNull { annotationCall ->
+            val annotationClass = annotationCall.symbol.owner.parentAsClass
+            if (!annotationClass.descriptor.isSerialInfoAnnotation) return@mapNotNull null
+
+            if (compilerContext.platform.isJvm()) {
+                val implClass = compilerContext.serialInfoImplJvmIrGenerator.getImplClass(annotationClass)
+                val ctor = implClass.constructors.singleOrNull { it.valueParameters.size == annotationCall.valueArgumentsCount }
+                    ?: error("No constructor args found for SerialInfo annotation Impl class: ${implClass.render()}")
+                irCall(ctor).apply {
+                    for (i in 0 until annotationCall.valueArgumentsCount) {
+                        val argument = annotationCall.getValueArgument(i)
+                            ?: annotationClass.primaryConstructor!!.valueParameters[i].defaultValue?.expression
+                        putValueArgument(i, argument!!.deepCopyWithVariables())
+                    }
+                }
+            } else {
+                annotationCall.deepCopyWithVariables()
+            }
+        }
+
     // Does not use sti and therefore does not perform encoder calls optimization
     fun IrBuilderWithScope.serializerTower(
         generator: SerializerIrGenerator,
@@ -920,8 +941,14 @@ interface IrBuilderExtension {
         val hasNewCtxSerCtor =
             serializerClassOriginal.classId == contextSerializerId && compilerContext.referenceConstructors(serializerClass.fqNameSafe)
                 .any { it.owner.valueParameters.size == 3 }
+        var needToCopyAnnotations = false
         when (serializerClassOriginal.classId) {
-            contextSerializerId, polymorphicSerializerId -> {
+            polymorphicSerializerId -> {
+                needToCopyAnnotations = true
+                args = listOf(classReference(kType))
+                typeArgs = listOf(thisIrType)
+            }
+            contextSerializerId -> {
                 args = listOf(classReference(kType))
                 typeArgs = listOf(thisIrType)
 
@@ -950,10 +977,12 @@ interface IrBuilderExtension {
                 }
             }
             objectSerializerId -> {
+                needToCopyAnnotations = true
                 args = listOf(irString(kType.serialName()), irGetObject(kType.toClassDescriptor!!))
                 typeArgs = listOf(thisIrType)
             }
             sealedSerializerId -> {
+                needToCopyAnnotations = true
                 args = mutableListOf<IrExpression>().apply {
                     add(irString(kType.serialName()))
                     add(classReference(kType))
@@ -1031,16 +1060,36 @@ interface IrBuilderExtension {
                 findSerializerConstructorForTypeArgumentsSerializers(serializerClass)
             ) { "Generated serializer does not have constructor with required number of arguments" }
         } else {
-            compilerContext.referenceConstructors(serializerClass.fqNameSafe).single { it.owner.isPrimary }
+            val constructors = compilerContext.referenceConstructors(serializerClass.fqNameSafe)
+            // search for new signature of polymorphic/sealed/contextual serializer
+            if (!needToCopyAnnotations) {
+                constructors.single { it.owner.isPrimary }
+            } else {
+                constructors.find { it.owner.lastArgumentIsAnnotationArray() } ?: run {
+                    // not found - we are using old serialization runtime without this feature
+                    needToCopyAnnotations = false
+                    constructors.single { it.owner.isPrimary }
+                }
+            }
         }
         // Return type should be correctly substituted
         assert(ctor.isBound)
         val ctorDecl = ctor.owner
+        if (needToCopyAnnotations) {
+            val classAnnotations = copyAnnotationsFrom(thisIrType.getClass()?.annotations.orEmpty())
+            args = args + createArrayOfExpression(compilerContext.builtIns.annotationType.toIrType(), classAnnotations)
+        }
+
         val typeParameters = ctorDecl.parentAsClass.typeParameters
         val substitutedReturnType = ctorDecl.returnType.substitute(typeParameters, typeArgs)
         return irInvoke(null, ctor, typeArguments = typeArgs, valueArguments = args, returnTypeHint = substitutedReturnType)
     }
 
+    private fun IrConstructor.lastArgumentIsAnnotationArray(): Boolean {
+        val lastArgType = valueParameters.lastOrNull()?.type
+        return (lastArgType != null && (lastArgType as? IrSimpleType)?.arguments?.firstOrNull()?.typeOrNull?.classFqName?.toString() == "kotlin.Annotation")
+    }
+
     private fun IrBuilderWithScope.wrapperClassReference(classType: KotlinType): IrClassReference {
         if (compilerContext.platform.isJvm()) {
             // "Byte::class" -> "java.lang.Byte::class"

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

@@ -156,25 +156,8 @@ open class SerializerIrGenerator(
         receiver: IrVariable,
         method: IrFunctionSymbol
     ) {
-        for (annotationCall in annotations) {
-            val annotationClass = annotationCall.symbol.owner.parentAsClass
-            if (!annotationClass.descriptor.isSerialInfoAnnotation) continue
-
-            val createAnnotation = if (compilerContext.platform.isJvm()) {
-                val implClass = serialInfoJvmGenerator.getImplClass(annotationClass)
-                val ctor = implClass.constructors.singleOrNull { it.valueParameters.size == annotationCall.valueArgumentsCount }
-                    ?: error("No constructor args found for SerialInfo annotation Impl class: ${implClass.render()}")
-                irCall(ctor).apply {
-                    for (i in 0 until annotationCall.valueArgumentsCount) {
-                        val argument = annotationCall.getValueArgument(i) ?: annotationClass.primaryConstructor!!.valueParameters[i].defaultValue?.expression
-                        putValueArgument(i, argument!!.deepCopyWithVariables())
-                    }
-                }
-            } else {
-                annotationCall.deepCopyWithVariables()
-            }
-
-            +irInvoke(irGet(receiver), method, createAnnotation)
+        copyAnnotationsFrom(annotations).forEach {
+            +irInvoke(irGet(receiver), method, it)
         }
     }
 

plugins/kotlin-serialization/kotlin-serialization-compiler/src/org/jetbrains/kotlinx/serialization/compiler/extensions/SerializationLoweringExtension.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2018 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
@@ -40,17 +40,22 @@ fun ClassLoweringPass.runOnFileInOrder(irFile: IrFile) {
     })
 }
 
-typealias SerializationPluginContext = IrPluginContext
+
+class SerializationPluginContext(baseContext: IrPluginContext, val metadataPlugin: SerializationDescriptorSerializerPlugin?) :
+    IrPluginContext by baseContext {
+    lateinit var serialInfoImplJvmIrGenerator: SerialInfoImplJvmIrGenerator
+}
 
 private class SerializerClassLowering(
-    val context: SerializationPluginContext,
-    val metadataPlugin: SerializationDescriptorSerializerPlugin?
+    baseContext: IrPluginContext,
+    metadataPlugin: SerializationDescriptorSerializerPlugin?
 ) : IrElementTransformerVoid(), ClassLoweringPass {
-    private val serialInfoJvmGenerator = SerialInfoImplJvmIrGenerator(context)
+    val context: SerializationPluginContext = SerializationPluginContext(baseContext, metadataPlugin)
+    private val serialInfoJvmGenerator = SerialInfoImplJvmIrGenerator(context).also { context.serialInfoImplJvmIrGenerator = it }
 
     override fun lower(irClass: IrClass) {
         SerializableIrGenerator.generate(irClass, context, context.bindingContext)
-        SerializerIrGenerator.generate(irClass, context, context.bindingContext, metadataPlugin, serialInfoJvmGenerator)
+        SerializerIrGenerator.generate(irClass, context, context.bindingContext, context.metadataPlugin, serialInfoJvmGenerator)
         SerializableCompanionIrGenerator.generate(irClass, context, context.bindingContext)
 
         if (context.platform.isJvm() && KSerializerDescriptorResolver.isSerialInfoImpl(irClass.descriptor)) {