Patch serializable class with addtional annotation

(which contains actual serializer) to make use of it with @AssociatedObjectKey

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

@@ -12,23 +12,36 @@ import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.annotations.Annotations
 import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.declarations.impl.*
+import org.jetbrains.kotlin.ir.declarations.impl.IrPropertyImpl
+import org.jetbrains.kotlin.ir.declarations.impl.IrTypeParameterImpl
+import org.jetbrains.kotlin.ir.declarations.impl.IrValueParameterImpl
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.*
 import org.jetbrains.kotlin.ir.symbols.*
-import org.jetbrains.kotlin.ir.types.*
+import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl
 import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
+import org.jetbrains.kotlin.ir.types.makeNotNull
+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.*
+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.findTypeSerializerOrContext
-import org.jetbrains.kotlinx.serialization.compiler.backend.jvm.*
+import org.jetbrains.kotlinx.serialization.compiler.backend.jvm.contextSerializerId
+import org.jetbrains.kotlinx.serialization.compiler.backend.jvm.enumSerializerId
+import org.jetbrains.kotlinx.serialization.compiler.backend.jvm.polymorphicSerializerId
+import org.jetbrains.kotlinx.serialization.compiler.backend.jvm.referenceArraySerializerId
 import org.jetbrains.kotlinx.serialization.compiler.resolve.*
 
 val BackendContext.externalSymbols: ReferenceSymbolTable get() = ir.symbols.externalSymbolTable
@@ -386,7 +399,7 @@ interface IrBuilderExtension {
         }
     }
 
-    fun IrBuilderWithScope.classReference(classType: KotlinType): IrClassReference {
+    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 =
@@ -400,6 +413,8 @@ interface IrBuilderExtension {
         )
     }
 
+    fun IrBuilderWithScope.classReference(classType: KotlinType): IrClassReference = createClassReference(classType, startOffset, endOffset)
+
     fun buildInitializersRemapping(irClass: IrClass): (IrField) -> IrExpression? {
         val original = irClass.constructors.singleOrNull { it.isPrimary }
             ?: throw IllegalStateException("Serializable class must have single primary constructor")

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

@@ -1,20 +1,24 @@
 package org.jetbrains.kotlinx.serialization.compiler.backend.ir
 
 import org.jetbrains.kotlin.backend.common.BackendContext
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.descriptors.findClassAcrossModuleDependencies
+import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
 import org.jetbrains.kotlin.ir.builders.irGet
 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.types.IrTypeArgument
+import org.jetbrains.kotlin.ir.expressions.impl.IrConstructorCallImpl
 import org.jetbrains.kotlin.ir.types.defaultType
-import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
 import org.jetbrains.kotlin.ir.util.SymbolTable
 import org.jetbrains.kotlin.ir.util.TypeTranslator
+import org.jetbrains.kotlin.ir.util.defaultType
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.descriptorUtil.module
-import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlinx.serialization.compiler.backend.common.SerializableCompanionCodegen
 import org.jetbrains.kotlinx.serialization.compiler.resolve.*
 
@@ -41,6 +45,38 @@ class SerializableCompanionIrGenerator(
         }
     }
 
+    private fun IrBuilderWithScope.patchSerializableClassWithMarkerAnnotation(serializer: ClassDescriptor) {
+        if (serializer.kind != ClassKind.OBJECT) {
+            return
+        }
+
+        val annotationMarkerClass = serializer.module.findClassAcrossModuleDependencies(
+            ClassId(
+                SerializationPackages.packageFqName,
+                Name.identifier(SerialEntityNames.ANNOTATION_MARKER_CLASS)
+            )
+        ) ?: return
+        val annotationCtor = requireNotNull(annotationMarkerClass.unsubstitutedPrimaryConstructor?.let {
+            compilerContext.externalSymbols.referenceConstructor(it)
+        })
+
+        val annotationType = compilerContext.externalSymbols.referenceClass(annotationMarkerClass).owner.defaultType
+        val irSerializableClass = compilerContext.externalSymbols.referenceClass(serializableDescriptor).owner
+        val annotationCtorCall = IrConstructorCallImpl.fromSymbolOwner(startOffset, endOffset, annotationType, annotationCtor).apply {
+            val serializerType = serializer.toSimpleType(false)
+            putValueArgument(
+                0,
+                createClassReference(
+                    serializerType,
+                    startOffset,
+                    endOffset
+                )
+            )
+        }
+
+        irSerializableClass.annotations.add(annotationCtorCall)
+    }
+
     override fun generateSerializerGetter(methodDescriptor: FunctionDescriptor) =
         irClass.contributeFunction(methodDescriptor, fromStubs = true) { getter ->
             val serializer = serializableDescriptor.classSerializer!!
@@ -66,6 +102,7 @@ class SerializableCompanionIrGenerator(
                     irInvoke(null, ctor, typeArgs, args, returnTypeHint = getter.returnType)
                 }
             }
+            patchSerializableClassWithMarkerAnnotation(serializer)
             +irReturn(expr)
         }
 }

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

@@ -48,6 +48,8 @@ object SerialEntityNames {
     const val DECODER_CLASS = "Decoder"
     const val STRUCTURE_DECODER_CLASS = "CompositeDecoder"
 
+    const val ANNOTATION_MARKER_CLASS = "SerializableWith"
+
     const val SERIAL_SAVER_CLASS = "SerializationStrategy"
     const val SERIAL_LOADER_CLASS = "DeserializationStrategy"