Insert correct dispatch receivers and type arguments for expression inside generated by plugin IR functions

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

@@ -18,9 +18,9 @@ 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.IrType
+import org.jetbrains.kotlin.ir.types.*
-import org.jetbrains.kotlin.ir.types.toKotlinType
+import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl
-import org.jetbrains.kotlin.ir.types.typeWith
+import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.js.resolve.diagnostics.findPsi
 import org.jetbrains.kotlin.name.ClassId
@@ -69,7 +69,7 @@ interface IrBuilderExtension {
         ) else compilerContext.externalSymbols.referenceSimpleFunction(descriptor).owner
         f.parent = this
         f.returnType = descriptor.returnType!!.toIrType()
-        if (!fromStubs) f.createParameterDeclarations(this.thisReceiver)
+        if (!fromStubs) f.createParameterDeclarations(this.thisReceiver!!)
         f.body = compilerContext.createIrBuilder(f.symbol).at(this).irBlockBody(this.startOffset, this.endOffset) { bodyGen(f) }
         this.addMember(f)
     }
@@ -87,7 +87,10 @@ interface IrBuilderExtension {
         ) else compilerContext.externalSymbols.referenceConstructor(descriptor).owner
         c.parent = this
         c.returnType = descriptor.returnType.toIrType()
-        if (!fromStubs) c.createParameterDeclarations()
+        if (!fromStubs) c.createParameterDeclarations(receiver = null)
+        if (c.typeParameters.isEmpty()) {
+            c.copyTypeParamsFromDescriptor()
+        }
         c.body = compilerContext.createIrBuilder(c.symbol).at(this).irBlockBody(this.startOffset, this.endOffset) { bodyGen(c) }
         this.addMember(c)
     }
@@ -237,9 +240,9 @@ interface IrBuilderExtension {
         irProperty.backingField =
             generatePropertyBackingField(propertyDescriptor, irProperty).apply { parent = propertyParent }
         val fieldSymbol = irProperty.backingField!!.symbol
-        irProperty.getter = propertyDescriptor.getter?.let { generatePropertyAccessor(it, fieldSymbol, ownerSymbol) }
+        irProperty.getter = propertyDescriptor.getter?.let { generatePropertyAccessor(it, fieldSymbol) }
             ?.apply { parent = propertyParent }
-        irProperty.setter = propertyDescriptor.setter?.let { generatePropertyAccessor(it, fieldSymbol, ownerSymbol) }
+        irProperty.setter = propertyDescriptor.setter?.let { generatePropertyAccessor(it, fieldSymbol) }
             ?.apply { parent = propertyParent }
         return irProperty
     }
@@ -259,18 +262,17 @@ interface IrBuilderExtension {
 
     fun generatePropertyAccessor(
         descriptor: PropertyAccessorDescriptor,
-        fieldSymbol: IrFieldSymbol,
+        fieldSymbol: IrFieldSymbol
-        ownerSymbol: IrValueSymbol
     ): IrSimpleFunction {
         // Declaration can also be called from user code. Since we lookup descriptor getter in externalSymbols
         // (see generateSave/generateLoad), seems it is correct approach to declare getter lazily there.
         val declaration = compilerContext.externalSymbols.referenceSimpleFunction(descriptor).owner
         return declaration.buildWithScope { irAccessor ->
-            irAccessor.createParameterDeclarations((ownerSymbol as IrValueParameterSymbol).owner) // todo: neat this
+            irAccessor.createParameterDeclarations(receiver = null)
             irAccessor.returnType = irAccessor.descriptor.returnType!!.toIrType()
             irAccessor.body = when (descriptor) {
-                is PropertyGetterDescriptor -> generateDefaultGetterBody(descriptor, irAccessor, ownerSymbol)
+                is PropertyGetterDescriptor -> generateDefaultGetterBody(descriptor, irAccessor)
-                is PropertySetterDescriptor -> generateDefaultSetterBody(descriptor, irAccessor, ownerSymbol)
+                is PropertySetterDescriptor -> generateDefaultSetterBody(descriptor, irAccessor)
                 else -> throw AssertionError("Should be getter or setter: $descriptor")
             }
         }
@@ -278,8 +280,7 @@ interface IrBuilderExtension {
 
     private fun generateDefaultGetterBody(
         getter: PropertyGetterDescriptor,
-        irAccessor: IrSimpleFunction,
+        irAccessor: IrSimpleFunction
-        ownerSymbol: IrValueSymbol
     ): IrBlockBody {
         val property = getter.correspondingProperty
 
@@ -287,7 +288,7 @@ interface IrBuilderExtension {
         val endOffset = irAccessor.endOffset
         val irBody = IrBlockBodyImpl(startOffset, endOffset)
 
-        val receiver = generateReceiverExpressionForFieldAccess(ownerSymbol, property)
+        val receiver = generateReceiverExpressionForFieldAccess(irAccessor.dispatchReceiverParameter!!.symbol, property)
 
         irBody.statements.add(
             IrReturnImpl(
@@ -306,8 +307,7 @@ interface IrBuilderExtension {
 
     private fun generateDefaultSetterBody(
         setter: PropertySetterDescriptor,
-        irAccessor: IrSimpleFunction,
+        irAccessor: IrSimpleFunction
-        ownerSymbol: IrValueSymbol
     ): IrBlockBody {
         val property = setter.correspondingProperty
 
@@ -315,7 +315,7 @@ interface IrBuilderExtension {
         val endOffset = irAccessor.endOffset
         val irBody = IrBlockBodyImpl(startOffset, endOffset)
 
-        val receiver = generateReceiverExpressionForFieldAccess(ownerSymbol, property)
+        val receiver = generateReceiverExpressionForFieldAccess(irAccessor.dispatchReceiverParameter!!.symbol, property)
 
         val irValueParameter = irAccessor.valueParameters.single()
         irBody.statements.add(
@@ -346,25 +346,28 @@ interface IrBuilderExtension {
         }
     }
 
-    // todo: delet zis
+    fun IrFunction.createParameterDeclarations(receiver: IrValueParameter?) {
-    fun IrFunction.createParameterDeclarations(receiver: IrValueParameter? = null) {
         fun ParameterDescriptor.irValueParameter() = IrValueParameterImpl(
             this@createParameterDeclarations.startOffset, this@createParameterDeclarations.endOffset,
             SERIALIZABLE_PLUGIN_ORIGIN,
             this,
             type.toIrType(),
-            (this as? ValueParameterDescriptor)?.varargElementType?.toIrType()
+            null
         ).also {
             it.parent = this@createParameterDeclarations
         }
 
-        dispatchReceiverParameter = receiver ?: (descriptor.dispatchReceiverParameter?.irValueParameter())
+        dispatchReceiverParameter = descriptor.dispatchReceiverParameter?.irValueParameter()
         extensionReceiverParameter = descriptor.extensionReceiverParameter?.irValueParameter()
 
         assert(valueParameters.isEmpty())
         descriptor.valueParameters.mapTo(valueParameters) { it.irValueParameter() }
 
         assert(typeParameters.isEmpty())
+        copyTypeParamsFromDescriptor()
+    }
+
+    fun IrFunction.copyTypeParamsFromDescriptor() {
         descriptor.typeParameters.mapTo(typeParameters) {
             IrTypeParameterImpl(
                 startOffset, endOffset,
@@ -376,8 +379,6 @@ interface IrBuilderExtension {
         }
     }
 
-
-
     fun IrBuilderWithScope.classReference(classType: KotlinType): IrClassReference {
         val clazz = classType.toClassDescriptor!!
         val kClass = clazz.module.findClassAcrossModuleDependencies(ClassId(FqName("kotlin.reflect"), Name.identifier("KClass")))!!
@@ -424,7 +425,7 @@ interface IrBuilderExtension {
     }
 
     // Does not use sti and therefore does not perform encoder calls optimization
-    fun IrBuilderWithScope.serializerTower(generator: SerializerIrGenerator, property: SerializableProperty): IrExpression? {
+    fun IrBuilderWithScope.serializerTower(generator: SerializerIrGenerator, dispatchReceiverParameter: IrValueParameter, property: SerializableProperty): IrExpression? {
         val nullableSerClass =
                 compilerContext.externalSymbols.referenceClass(property.module.getClassFromInternalSerializationPackage(SpecialBuiltins.nullableSerializer))
         val serializer =
@@ -435,19 +436,35 @@ interface IrBuilderExtension {
                                 property.descriptor.annotations,
                                 property.descriptor.findPsi()
                         ) else null
-        return serializerInstance(generator, generator.serializableDescriptor, serializer, property.module, property.type, property.genericIndex)
+        return serializerInstance(generator, dispatchReceiverParameter, generator.serializableDescriptor, serializer, property.module, property.type, genericIndex = property.genericIndex)
-                ?.let { expr -> wrapWithNullableSerializerIfNeeded(property.type, expr, nullableSerClass) }
+                ?.let { expr -> wrapWithNullableSerializerIfNeeded(property.module, property.type, expr, nullableSerClass) }
     }
 
-    private fun IrBuilderWithScope.wrapWithNullableSerializerIfNeeded(type: KotlinType, expression: IrExpression, nullableSerializerClass: IrClassSymbol): IrExpression {
+    private fun IrBuilderWithScope.wrapWithNullableSerializerIfNeeded(module: ModuleDescriptor, type: KotlinType, expression: IrExpression, nullableSerializerClass: IrClassSymbol): IrExpression {
         return if (type.isMarkedNullable)
-            irInvoke(null, nullableSerializerClass.constructors.toList()[0], listOf(type.makeNotNullable().toIrType()), listOf(expression))
+            irInvoke(null, nullableSerializerClass.constructors.toList()[0],
+                typeArguments  = listOf(type.makeNotNullable().toIrType()),
+                valueArguments = listOf(expression),
+                returnTypeHint = wrapIrTypeIntoKSerializerIrType(module, type.toIrType())
+            )
         else
             expression
     }
 
+
+    fun wrapIrTypeIntoKSerializerIrType(module: ModuleDescriptor, type: IrType): IrType {
+        val kSerClass =
+            compilerContext.externalSymbols.referenceClass(module.getClassFromSerializationPackage(SerialEntityNames.KSERIALIZER_CLASS))
+        return IrSimpleTypeImpl(
+            kSerClass, hasQuestionMark = false, arguments = listOf(
+                makeTypeProjection(type, Variance.INVARIANT)
+            ), annotations = emptyList()
+        )
+    }
+
     fun IrBuilderWithScope.serializerInstance(
         enclosingGenerator: SerializerIrGenerator,
+        dispatchReceiverParameter: IrValueParameter,
         serializableDescriptor: ClassDescriptor,
         serializerClassOriginal: ClassDescriptor?,
         module: ModuleDescriptor,
@@ -460,17 +477,22 @@ interface IrBuilderExtension {
             if (genericIndex == null) return null
             val thiz = enclosingGenerator.irClass.thisReceiver!!
             val prop = enclosingGenerator.localSerializersFieldsDescriptors[genericIndex]
-            return irGetField(irGet(thiz), compilerContext.localSymbolTable.referenceField(prop).owner)
+            return irGetField(irGet(dispatchReceiverParameter), compilerContext.localSymbolTable.referenceField(prop).owner)
         }
         if (serializerClassOriginal.kind == ClassKind.OBJECT) {
             return irGetObject(serializerClassOriginal)
         } else {
             var serializerClass = serializerClassOriginal
-            var args: List<IrExpression> = when (serializerClassOriginal.classId) {
+            var args: List<IrExpression>
-                contextSerializerId, polymorphicSerializerId -> listOf(classReference(kType))
+            var typeArgs: List<IrType?>
+            when (serializerClassOriginal.classId) {
+                contextSerializerId, polymorphicSerializerId -> {
+                    args = listOf(classReference(kType))
+                    typeArgs = listOf(kType.toIrType())
+                }
                 enumSerializerId -> {
                     serializerClass = serializableDescriptor.getClassFromInternalSerializationPackage("CommonEnumSerializer")
-                    kType.toClassDescriptor!!.let { enumDesc ->
+                    args = kType.toClassDescriptor!!.let { enumDesc ->
                         listOf(
                             irString(enumDesc.name.toString()),
                             irCall(findEnumValuesMethod(enumDesc)),
@@ -480,18 +502,25 @@ interface IrBuilderExtension {
                             )
                         )
                     }
+                    typeArgs = listOf(kType.toIrType())
                 }
-                else -> kType.arguments.map {
+                else -> {
-                    val argSer = enclosingGenerator.findTypeSerializerOrContext(module, it.type, sourceElement = serializerClassOriginal.findPsi())
+                    args = kType.arguments.map {
-                    val expr = serializerInstance(enclosingGenerator, serializableDescriptor, argSer, module, it.type, it.type.genericIndex)
+                        val argSer = enclosingGenerator.findTypeSerializerOrContext(module, it.type, sourceElement = serializerClassOriginal.findPsi())
-                        ?: return null
+                        val expr = serializerInstance(enclosingGenerator, dispatchReceiverParameter, serializableDescriptor, argSer, module, it.type, it.type.genericIndex)
-                    wrapWithNullableSerializerIfNeeded(it.type, expr, nullableSerClass)
+                            ?: return null
+                        wrapWithNullableSerializerIfNeeded(module, it.type, expr, nullableSerClass)
+                    }
+                    typeArgs = kType.arguments.map { it.type.toIrType() }
                 }
-            }
-            if (serializerClassOriginal.classId == referenceArraySerializerId)
-                args = listOf(classReference(kType.arguments[0].type)) + args
 
-            val typeArgs = kType.arguments.map { it.type.toIrType() }
+            }
+            if (serializerClassOriginal.classId == referenceArraySerializerId) {
+                args = listOf(classReference(kType.arguments[0].type)) + args
+                typeArgs = listOf(typeArgs[0].makeNotNull()) + typeArgs
+            }
+
+
             val serializable = getSerializableClassDescriptorBySerializer(serializerClass)
             val ctor = if (serializable?.declaredTypeParameters?.isNotEmpty() == true) {
                 requireNotNull(
@@ -501,7 +530,9 @@ interface IrBuilderExtension {
             } else {
                 compilerContext.externalSymbols.referenceConstructor(serializerClass.unsubstitutedPrimaryConstructor!!)
             }
-            return irInvoke(null, ctor, typeArgs, args)
+            val propertyTypeHint = kType.toIrType()
+            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/ir/SerializableCompanionIrGenerator.kt

@@ -7,10 +7,14 @@ 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.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.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.*
 
@@ -47,8 +51,9 @@ class SerializableCompanionIrGenerator(
                     irInvoke(
                         null,
                         compilerContext.externalSymbols.referenceConstructor(serializer.unsubstitutedPrimaryConstructor!!),
-                        listOf(serializableType.toIrType()),
+                        typeArguments = listOf(serializableType.toIrType()),
-                        listOf(classReference(serializableType))
+                        valueArguments = listOf(classReference(serializableType)),
+                        returnTypeHint = getter.returnType
                     )
                 }
                 else -> {
@@ -56,8 +61,9 @@ class SerializableCompanionIrGenerator(
                         KSerializerDescriptorResolver.findSerializerConstructorForTypeArgumentsSerializers(serializer)
                     ) { "Generated serializer does not have constructor with required number of arguments" }
                     val ctor = compilerContext.externalSymbols.referenceConstructor(desc)
+                    val typeArgs = getter.typeParameters.map { it.defaultType }
                     val args: List<IrExpression> = getter.valueParameters.map { irGet(it) }
-                    irInvoke(null, ctor, *args.toTypedArray())
+                    irInvoke(null, ctor, typeArgs, args, returnTypeHint = getter.returnType)
                 }
             }
             +irReturn(expr)

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

@@ -152,7 +152,11 @@ class SerializerIrGenerator(val irClass: IrClass, override val compilerContext:
                 endOffset,
                 compilerContext.irBuiltIns.unitType,
                 primaryCtor.symbol
-            )
+            ).apply {
+                ctor.typeParameters.forEachIndexed { index, irTypeParameter ->
+                    putTypeArgument(index, irTypeParameter.defaultType)
+                }
+            }
 
             // store type arguments serializers in fields
             val thisAsReceiverParameter = irClass.thisReceiver!!
@@ -163,15 +167,17 @@ class SerializerIrGenerator(val irClass: IrClass, override val compilerContext:
                     localSerializersFieldsDescriptors[index]
                 ), localSerial.owner, irGet(param))
             }
+
+
         }
 
     override fun generateChildSerializersGetter(function: FunctionDescriptor) = irClass.contributeFunction(function) { irFun ->
         val allSerializers = serializableProperties.map { requireNotNull(
-                serializerTower(this@SerializerIrGenerator, it)) { "Property ${it.name} must have a serializer" }
+                serializerTower(this@SerializerIrGenerator, irFun.dispatchReceiverParameter!!, it)) { "Property ${it.name} must have a serializer" }
         }
 
         val kSer = serializableDescriptor.module.getClassFromSerializationPackage(KSERIALIZER_NAME.identifier)
-        val kSerType = compilerContext.externalSymbols.referenceClass(kSer).owner.defaultType
+        val kSerType = ((irFun.returnType as IrSimpleType).arguments.first() as IrTypeProjection).type
         val array = createArrayOfExpression(kSerType, allSerializers)
         +irReturn(array)
     }
@@ -222,11 +228,12 @@ class SerializerIrGenerator(val irClass: IrClass, override val compilerContext:
             val sti = getSerialTypeInfo(property)
             val innerSerial = serializerInstance(
                 this@SerializerIrGenerator,
+                saveFunc.dispatchReceiverParameter!!,
                 serializableDescriptor,
                 sti.serializer,
                 property.module,
                 property.type,
-                property.genericIndex
+                genericIndex = property.genericIndex
             )
             val (writeFunc, args: List<IrExpression>) = if (innerSerial == null) {
                 val f =
@@ -351,11 +358,12 @@ class SerializerIrGenerator(val irClass: IrClass, override val compilerContext:
                                 val sti = getSerialTypeInfo(property)
                                 val innerSerial = serializerInstance(
                                     this@SerializerIrGenerator,
+                                    loadFunc.dispatchReceiverParameter!!,
                                     serializableDescriptor,
                                     sti.serializer,
                                     property.module,
                                     property.type,
-                                    property.genericIndex
+                                    genericIndex = property.genericIndex
                                 )
                                 // todo: update
                                 val decodeFuncToCall =
@@ -372,7 +380,7 @@ class SerializerIrGenerator(val irClass: IrClass, override val compilerContext:
                                 // local$i = localInput.decode...(...)
                                 +irSetVar(
                                     localProps[index].symbol,
-                                    irInvoke(localInput.get(), decodeFuncToCall, typeArgs, args)
+                                    irInvoke(localInput.get(), decodeFuncToCall, typeArgs, args, returnTypeHint = property.type.toIrType())
                                 )
                                 // bitMask[i] |= 1 << x
                                 val bitPos = 1 shl (index % 32)
@@ -412,6 +420,7 @@ class SerializerIrGenerator(val irClass: IrClass, override val compilerContext:
 
         // todo: set properties in external deserialization
         var args: List<IrExpression> = localProps.map { it.get() }
+        val typeArgs = (loadFunc.returnType as IrSimpleType).arguments.map { (it as IrTypeProjection).type }
         val ctor: IrConstructorSymbol = if (serializableDescriptor.isInternalSerializable) {
             args = bitMasks.map { irGet(it) } + args + irNull()
             serializableIrClass.serializableSyntheticConstructor()
@@ -419,7 +428,7 @@ class SerializerIrGenerator(val irClass: IrClass, override val compilerContext:
             compilerContext.externalSymbols.referenceConstructor(serializableDescriptor.unsubstitutedPrimaryConstructor!!)
         }
 
-        +irReturn(irInvoke(null, ctor, *args.toTypedArray()))
+        +irReturn(irInvoke(null, ctor, typeArgs, args))
     }
 
     companion object {