IrTypes: IrFunction.returnType can depend on type parameters

Function return type can depend on function type parameters and should
be initialized with type parameters in scope.

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/ClassGenerator.kt

@@ -207,10 +207,13 @@ class ClassGenerator(
         context.symbolTable.declareSimpleFunctionWithOverrides(
             irDelegate.startOffset, irDelegate.endOffset,
             IrDeclarationOrigin.DELEGATED_MEMBER,
-            delegated,
-            delegated.returnType!!.toIrType()
+            delegated
         ).buildWithScope { irFunction ->
             FunctionGenerator(declarationGenerator).generateSyntheticFunctionParameterDeclarations(irFunction)
+
+            // TODO could possibly refer to scoped type parameters for property accessors
+            irFunction.returnType = delegated.returnType!!.toIrType()
+
             irFunction.body = generateDelegateFunctionBody(irDelegate, delegated, overridden, irFunction)
         }
 

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/DataClassMembersGenerator.kt

@@ -55,8 +55,10 @@ class DataClassMembersGenerator(
     private fun declareSimpleFunction(startOffset: Int, endOffset: Int, origin: IrDeclarationOrigin, function: FunctionDescriptor) =
         context.symbolTable.declareSimpleFunctionWithOverrides(
             startOffset, endOffset, origin,
-            function, function.returnType!!.toIrType()
-        )
+            function
+        ).apply {
+            returnType = function.returnType!!.toIrType()
+        }
 
     private inner class MemberFunctionBuilder(
         val irClass: IrClass,

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/DelegatedPropertyGenerator.kt

@@ -100,9 +100,9 @@ class DelegatedPropertyGenerator(declarationGenerator: DeclarationGenerator) : D
         context.symbolTable.declareSimpleFunctionWithOverrides(
             ktDelegate.startOffset, ktDelegate.endOffset,
             IrDeclarationOrigin.DELEGATED_PROPERTY_ACCESSOR,
-            accessorDescriptor, accessorDescriptor.returnType!!.toIrType()
+            accessorDescriptor
         ).buildWithScope { irAccessor ->
-            FunctionGenerator(declarationGenerator).generateFunctionParameterDeclarations(irAccessor, ktProperty, null)
+            FunctionGenerator(declarationGenerator).generateFunctionParameterDeclarationsAndReturnType(irAccessor, ktProperty, null)
             irAccessor.body = generateBody(irAccessor)
         }
 
@@ -338,9 +338,9 @@ class DelegatedPropertyGenerator(declarationGenerator: DeclarationGenerator) : D
         context.symbolTable.declareSimpleFunctionWithOverrides(
             ktDelegate.startOffset, ktDelegate.endOffset,
             IrDeclarationOrigin.DELEGATED_PROPERTY_ACCESSOR,
-            getterDescriptor, getterDescriptor.returnType!!.toIrType()
+            getterDescriptor
         ).buildWithScope { irAccessor ->
-            FunctionGenerator(declarationGenerator).generateFunctionParameterDeclarations(irAccessor, ktDelegate, null)
+            FunctionGenerator(declarationGenerator).generateFunctionParameterDeclarationsAndReturnType(irAccessor, ktDelegate, null)
             irAccessor.body = generateBody(irAccessor)
         }
 

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/EnumClassMembersGenerator.kt

@@ -42,9 +42,9 @@ class EnumClassMembersGenerator(declarationGenerator: DeclarationGenerator) : De
             context.symbolTable.declareSimpleFunctionWithOverrides(
                 irClass.startOffset, irClass.endOffset,
                 IrDeclarationOrigin.ENUM_CLASS_SPECIAL_MEMBER,
-                valuesFunction, valuesFunction.returnType!!.toIrType()
+                valuesFunction
             ).also { irFunction ->
-                FunctionGenerator(declarationGenerator).generateFunctionParameterDeclarations(irFunction, null, null)
+                FunctionGenerator(declarationGenerator).generateFunctionParameterDeclarationsAndReturnType(irFunction, null, null)
                 irFunction.body = IrSyntheticBodyImpl(irClass.startOffset, irClass.endOffset, IrSyntheticBodyKind.ENUM_VALUES)
             }
         )
@@ -61,9 +61,9 @@ class EnumClassMembersGenerator(declarationGenerator: DeclarationGenerator) : De
             context.symbolTable.declareSimpleFunctionWithOverrides(
                 UNDEFINED_OFFSET, UNDEFINED_OFFSET,
                 IrDeclarationOrigin.ENUM_CLASS_SPECIAL_MEMBER,
-                valueOfFunction, valueOfFunction.returnType!!.toIrType()
+                valueOfFunction
             ).also { irFunction ->
-                FunctionGenerator(declarationGenerator).generateFunctionParameterDeclarations(irFunction, null, null)
+                FunctionGenerator(declarationGenerator).generateFunctionParameterDeclarationsAndReturnType(irFunction, null, null)
                 irFunction.body = IrSyntheticBodyImpl(irClass.startOffset, irClass.endOffset, IrSyntheticBodyKind.ENUM_VALUEOF)
             }
         )

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/FunctionGenerator.kt

@@ -59,7 +59,7 @@ class FunctionGenerator(declarationGenerator: DeclarationGenerator) : Declaratio
 
     fun generateFakeOverrideFunction(functionDescriptor: FunctionDescriptor, ktElement: KtElement): IrSimpleFunction =
         declareSimpleFunctionInner(functionDescriptor, ktElement, IrDeclarationOrigin.FAKE_OVERRIDE).buildWithScope { irFunction ->
-            generateFunctionParameterDeclarations(irFunction, ktElement, null)
+            generateFunctionParameterDeclarationsAndReturnType(irFunction, ktElement, null)
         }
 
     private inline fun declareSimpleFunction(
@@ -70,7 +70,7 @@ class FunctionGenerator(declarationGenerator: DeclarationGenerator) : Declaratio
         generateBody: BodyGenerator.() -> IrBody?
     ): IrSimpleFunction =
         declareSimpleFunctionInner(descriptor, ktFunction, origin).buildWithScope { irFunction ->
-            generateFunctionParameterDeclarations(irFunction, ktFunction, ktReceiver)
+            generateFunctionParameterDeclarationsAndReturnType(irFunction, ktFunction, ktReceiver)
             irFunction.body = createBodyGenerator(irFunction.symbol).generateBody()
         }
 
@@ -81,15 +81,16 @@ class FunctionGenerator(declarationGenerator: DeclarationGenerator) : Declaratio
     ): IrSimpleFunction =
         context.symbolTable.declareSimpleFunctionWithOverrides(
             ktElement.startOffset, ktElement.endOffset, origin,
-            descriptor, descriptor.returnType!!.toIrType()
+            descriptor
         )
 
-    fun generateFunctionParameterDeclarations(
+    fun generateFunctionParameterDeclarationsAndReturnType(
         irFunction: IrFunction,
         ktParameterOwner: KtElement?,
         ktReceiverParameterElement: KtElement?
     ) {
         declarationGenerator.generateScopedTypeParameterDeclarations(irFunction, irFunction.descriptor.typeParameters)
+        irFunction.returnType = irFunction.descriptor.returnType!!.toIrType()
         generateValueParameterDeclarations(irFunction, ktParameterOwner, ktReceiverParameterElement)
     }
 
@@ -104,7 +105,7 @@ class FunctionGenerator(declarationGenerator: DeclarationGenerator) : Declaratio
             if (ktAccessor != null) IrDeclarationOrigin.DEFINED else IrDeclarationOrigin.DEFAULT_PROPERTY_ACCESSOR
         ).buildWithScope { irAccessor ->
             declarationGenerator.generateScopedTypeParameterDeclarations(irAccessor, descriptor.correspondingProperty.typeParameters)
-            generateFunctionParameterDeclarations(irAccessor, ktAccessor ?: ktProperty, ktProperty.receiverTypeReference)
+            generateFunctionParameterDeclarationsAndReturnType(irAccessor, ktAccessor ?: ktProperty, ktProperty.receiverTypeReference)
             val ktBodyExpression = ktAccessor?.bodyExpression
             irAccessor.body =
                     if (ktBodyExpression != null)
@@ -236,10 +237,11 @@ class FunctionGenerator(declarationGenerator: DeclarationGenerator) : Declaratio
     ): IrConstructor =
         context.symbolTable.declareConstructor(
             ktConstructorElement.startOffset, ktConstructorElement.endOffset, IrDeclarationOrigin.DEFINED,
-            constructorDescriptor, constructorDescriptor.returnType.toIrType()
+            constructorDescriptor
         ).buildWithScope { irConstructor ->
             generateValueParameterDeclarations(irConstructor, ktParametersElement, null)
             irConstructor.body = createBodyGenerator(irConstructor.symbol).generateBody()
+            irConstructor.returnType = constructorDescriptor.returnType.toIrType()
         }
 
     fun generateSyntheticFunctionParameterDeclarations(irFunction: IrFunction) {

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/builders/IrMemberFunctionBuilder.kt

@@ -42,9 +42,10 @@ class IrMemberFunctionBuilder(
     lateinit var irFunction: IrFunction
 
     inline fun addToClass(body: IrMemberFunctionBuilder.(IrFunction) -> Unit): IrFunction {
-        irFunction = IrFunctionImpl(startOffset, endOffset, origin, function, returnType)
+        irFunction = IrFunctionImpl(startOffset, endOffset, origin, function)
         body(irFunction)
         irFunction.body = doBuild()
+        irFunction.returnType = returnType
         irClass.declarations.add(irFunction)
         return irFunction
     }

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/IrFunction.kt

@@ -33,7 +33,7 @@ interface IrFunction : IrDeclaration, IrTypeParametersContainer, IrSymbolOwner,
     val visibility: Visibility
     val isInline: Boolean // NB: there's an inline constructor for Array and each primitive array class
     val isExternal: Boolean
-    val returnType: IrType
+    var returnType: IrType
 
     var dispatchReceiverParameter: IrValueParameter?
     var extensionReceiverParameter: IrValueParameter?

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/impl/IrConstructorImpl.kt

@@ -35,14 +35,13 @@ class IrConstructorImpl(
     override val symbol: IrConstructorSymbol,
     name: Name,
     visibility: Visibility,
-    returnType: IrType,
     isInline: Boolean,
     isExternal: Boolean,
     override val isPrimary: Boolean
 ) :
     IrFunctionBase(
         startOffset, endOffset, origin, name,
-        visibility, isInline, isExternal, returnType
+        visibility, isInline, isExternal
     ),
     IrConstructor {
 
@@ -51,13 +50,11 @@ class IrConstructorImpl(
         endOffset: Int,
         origin: IrDeclarationOrigin,
         symbol: IrConstructorSymbol,
-        returnType: IrType,
         body: IrBody? = null
     ) : this(
         startOffset, endOffset, origin, symbol,
         symbol.descriptor.name,
         symbol.descriptor.visibility,
-        returnType,
         symbol.descriptor.isInline,
         symbol.descriptor.isEffectivelyExternal(),
         symbol.descriptor.isPrimary
@@ -69,19 +66,17 @@ class IrConstructorImpl(
         startOffset: Int,
         endOffset: Int,
         origin: IrDeclarationOrigin,
-        descriptor: ClassConstructorDescriptor,
-        returnType: IrType
-    ) : this(startOffset, endOffset, origin, IrConstructorSymbolImpl(descriptor), returnType)
+        descriptor: ClassConstructorDescriptor
+    ) : this(startOffset, endOffset, origin, IrConstructorSymbolImpl(descriptor))
 
-    @Deprecated("Let use constructor which takes symbol instead of descriptor")
+    @Deprecated("Use constructor which takes symbol instead of descriptor")
     constructor(
         startOffset: Int,
         endOffset: Int,
         origin: IrDeclarationOrigin,
         descriptor: ClassConstructorDescriptor,
-        returnType: IrType,
         body: IrBody?
-    ) : this(startOffset, endOffset, origin, descriptor, returnType) {
+    ) : this(startOffset, endOffset, origin, descriptor) {
         this.body = body
     }
 

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/impl/IrFunctionBase.kt

@@ -36,9 +36,9 @@ abstract class IrFunctionBase(
     override val name: Name,
     override val visibility: Visibility,
     override val isInline: Boolean,
-    override val isExternal: Boolean,
-    override val returnType: IrType
-) : IrDeclarationBase(startOffset, endOffset, origin),
+    override val isExternal: Boolean
+) :
+    IrDeclarationBase(startOffset, endOffset, origin),
     IrFunction {
 
     override val typeParameters: MutableList<IrTypeParameter> = SmartList()
@@ -49,6 +49,8 @@ abstract class IrFunctionBase(
 
     final override var body: IrBody? = null
 
+    final override lateinit var returnType: IrType
+
     override fun <D> acceptChildren(visitor: IrElementVisitor<Unit, D>, data: D) {
         typeParameters.forEach { it.accept(visitor, data) }
 

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/impl/IrFunctionImpl.kt

@@ -38,27 +38,24 @@ class IrFunctionImpl(
     name: Name,
     visibility: Visibility,
     override val modality: Modality,
-    returnType: IrType,
     isInline: Boolean,
     isExternal: Boolean,
     override val isTailrec: Boolean,
     override val isSuspend: Boolean
 ) :
-    IrFunctionBase(startOffset, endOffset, origin, name, visibility, isInline, isExternal, returnType),
+    IrFunctionBase(startOffset, endOffset, origin, name, visibility, isInline, isExternal),
     IrSimpleFunction {
 
     constructor(
         startOffset: Int,
         endOffset: Int,
         origin: IrDeclarationOrigin,
-        symbol: IrSimpleFunctionSymbol,
-        returnType: IrType
+        symbol: IrSimpleFunctionSymbol
     ) : this(
         startOffset, endOffset, origin, symbol,
         symbol.descriptor.name,
         symbol.descriptor.visibility,
         symbol.descriptor.modality,
-        returnType,
         symbol.descriptor.isInline,
         symbol.descriptor.isExternal,
         symbol.descriptor.isTailrec,
@@ -75,12 +72,10 @@ class IrFunctionImpl(
         startOffset: Int,
         endOffset: Int,
         origin: IrDeclarationOrigin,
-        descriptor: FunctionDescriptor,
-        returnType: IrType
+        descriptor: FunctionDescriptor
     ) : this(
         startOffset, endOffset, origin,
-        IrSimpleFunctionSymbolImpl(descriptor),
-        returnType
+        IrSimpleFunctionSymbolImpl(descriptor)
     )
 
     constructor(
@@ -88,9 +83,8 @@ class IrFunctionImpl(
         endOffset: Int,
         origin: IrDeclarationOrigin,
         descriptor: FunctionDescriptor,
-        returnType: IrType,
         body: IrBody?
-    ) : this(startOffset, endOffset, origin, descriptor, returnType) {
+    ) : this(startOffset, endOffset, origin, descriptor) {
         this.body = body
     }
 

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/DeclarationStubGenerator.kt

@@ -85,17 +85,23 @@ class DeclarationStubGenerator(
             } else {
                 origin
             },
-            descriptor.original,
-            descriptor.returnType!!.toIrType()
+            descriptor.original
         ).also { irFunction ->
             generateTypeParameterStubs(descriptor.propertyIfAccessor.typeParameters, irFunction)
+
+            // TODO make sure that type parameters have proper scopes
+            irFunction.returnType = descriptor.returnType!!.toIrType()
+
             generateValueParametersStubs(irFunction)
         }
 
     private fun generateConstructorStub(descriptor: ClassConstructorDescriptor): IrConstructor =
         symbolTable.declareConstructor(
-            UNDEFINED_OFFSET, UNDEFINED_OFFSET, origin, descriptor.original, descriptor.returnType.toIrType()
+            UNDEFINED_OFFSET, UNDEFINED_OFFSET, origin, descriptor.original
         ).also { irConstructor ->
+            // So far, constructors in Kotlin can't have type parameters of their own.
+            irConstructor.returnType = descriptor.returnType.toIrType()
+
             generateValueParametersStubs(irConstructor)
         }
 

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/DeepCopyIrTree.kt

@@ -116,7 +116,6 @@ open class DeepCopyIrTree : IrElementTransformerVoid() {
             declaration.startOffset, declaration.endOffset,
             mapDeclarationOrigin(declaration.origin),
             mapFunctionDeclaration(declaration.descriptor),
-            declaration.returnType, // TODO
             declaration.body?.transform()
         ).transformParameters(declaration).apply {
             transformAnnotations(declaration)
@@ -127,6 +126,7 @@ open class DeepCopyIrTree : IrElementTransformerVoid() {
                 else
                     IrSimpleFunctionSymbolImpl(overriddenDescriptor.original)
             }
+            returnType = declaration.returnType // TODO
         }
 
     override fun visitConstructor(declaration: IrConstructor): IrConstructor =
@@ -134,10 +134,10 @@ open class DeepCopyIrTree : IrElementTransformerVoid() {
             declaration.startOffset, declaration.endOffset,
             mapDeclarationOrigin(declaration.origin),
             mapConstructorDeclaration(declaration.descriptor),
-            declaration.returnType,
             declaration.body?.transform()
         ).transformParameters(declaration).apply {
             transformAnnotations(declaration)
+            returnType = declaration.returnType // TODO
         }
 
     protected fun <T : IrTypeParametersContainer> T.transformTypeParameters(

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/DeepCopyIrTreeWithSymbols.kt

@@ -120,8 +120,7 @@ open class DeepCopyIrTreeWithSymbols(
         IrFunctionImpl(
             declaration.startOffset, declaration.endOffset,
             mapDeclarationOrigin(declaration.origin),
-            symbolRemapper.getDeclaredFunction(declaration.symbol),
-            declaration.returnType.remapType()
+            symbolRemapper.getDeclaredFunction(declaration.symbol)
         ).apply {
             declaration.overriddenSymbols.mapTo(overriddenSymbols) {
                 symbolRemapper.getReferencedFunction(it) as IrSimpleFunctionSymbol
@@ -133,8 +132,7 @@ open class DeepCopyIrTreeWithSymbols(
         IrConstructorImpl(
             declaration.startOffset, declaration.endOffset,
             mapDeclarationOrigin(declaration.origin),
-            symbolRemapper.getDeclaredConstructor(declaration.symbol),
-            declaration.returnType.remapType()
+            symbolRemapper.getDeclaredConstructor(declaration.symbol)
         ).apply {
             transformFunctionChildren(declaration)
         }
@@ -143,10 +141,13 @@ open class DeepCopyIrTreeWithSymbols(
         apply {
             transformAnnotations(declaration)
             declaration.typeParameters.transformTo(typeParameters)
-            dispatchReceiverParameter = declaration.dispatchReceiverParameter?.transform()
-            extensionReceiverParameter = declaration.extensionReceiverParameter?.transform()
-            declaration.valueParameters.transformTo(valueParameters)
-            body = declaration.body?.transform()
+            typeRemapper.withinScope(this) {
+                dispatchReceiverParameter = declaration.dispatchReceiverParameter?.transform()
+                extensionReceiverParameter = declaration.extensionReceiverParameter?.transform()
+                returnType = typeRemapper.remapType(declaration.returnType)
+                declaration.valueParameters.transformTo(valueParameters)
+                body = declaration.body?.transform()
+            }
         }
 
     private fun IrAnnotationContainer.transformAnnotations(declaration: IrAnnotationContainer) {

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/DeepCopyTypeRemapper.kt

@@ -5,12 +5,21 @@
 
 package org.jetbrains.kotlin.ir.util
 
+import org.jetbrains.kotlin.ir.declarations.IrTypeParametersContainer
 import org.jetbrains.kotlin.ir.types.IrType
 
 class DeepCopyTypeRemapper(
     private val symbolRemapper: SymbolRemapper
 ) : TypeRemapper {
 
+    override fun enterScope(irTypeParametersContainer: IrTypeParametersContainer) {
+        // TODO
+    }
+
+    override fun leaveScope() {
+        // TODO
+    }
+
     override fun remapType(type: IrType): IrType = type // TODO
 
 }

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/SymbolTable.kt

@@ -190,13 +190,12 @@ class SymbolTable {
         startOffset: Int,
         endOffset: Int,
         origin: IrDeclarationOrigin,
-        descriptor: ClassConstructorDescriptor,
-        returnType: IrType
+        descriptor: ClassConstructorDescriptor
     ): IrConstructor =
         constructorSymbolTable.declare(
             descriptor,
             { IrConstructorSymbolImpl(descriptor) },
-            { IrConstructorImpl(startOffset, endOffset, origin, it, returnType) }
+            { IrConstructorImpl(startOffset, endOffset, origin, it) }
         )
 
     fun referenceConstructor(descriptor: ClassConstructorDescriptor) =
@@ -250,13 +249,12 @@ class SymbolTable {
         startOffset: Int,
         endOffset: Int,
         origin: IrDeclarationOrigin,
-        descriptor: FunctionDescriptor,
-        returnType: IrType
+        descriptor: FunctionDescriptor
     ): IrSimpleFunction =
         simpleFunctionSymbolTable.declare(
             descriptor,
             { IrSimpleFunctionSymbolImpl(descriptor) },
-            { IrFunctionImpl(startOffset, endOffset, origin, it, returnType) }
+            { IrFunctionImpl(startOffset, endOffset, origin, it) }
         )
 
     fun referenceSimpleFunction(descriptor: FunctionDescriptor) =

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/TypeRemapper.kt

@@ -5,8 +5,18 @@
 
 package org.jetbrains.kotlin.ir.util
 
+import org.jetbrains.kotlin.ir.declarations.IrTypeParametersContainer
 import org.jetbrains.kotlin.ir.types.IrType
 
 interface TypeRemapper {
+    fun enterScope(irTypeParametersContainer: IrTypeParametersContainer)
     fun remapType(type: IrType): IrType
+    fun leaveScope()
+}
+
+inline fun <T> TypeRemapper.withinScope(irTypeParametersContainer: IrTypeParametersContainer, fn: () -> T): T {
+    enterScope(irTypeParametersContainer)
+    val result = fn()
+    leaveScope()
+    return result
 }

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/overrides.kt

@@ -14,10 +14,9 @@ fun SymbolTable.declareSimpleFunctionWithOverrides(
     startOffset: Int,
     endOffset: Int,
     origin: IrDeclarationOrigin,
-    descriptor: FunctionDescriptor,
-    returnType: IrType
+    descriptor: FunctionDescriptor
 ) =
-    declareSimpleFunction(startOffset, endOffset, origin, descriptor, returnType).also { declaration ->
+    declareSimpleFunction(startOffset, endOffset, origin, descriptor).also { declaration ->
         generateOverriddenSymbols(declaration, this)
     }