Implemented inheritance of lambdas from SuspendFunction

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/common/descriptors/DescriptorUtils.kt

@@ -18,6 +18,12 @@ package org.jetbrains.kotlin.backend.common.descriptors
 
 import org.jetbrains.kotlin.coroutines.isSuspendLambda
 import org.jetbrains.kotlin.descriptors.*
+import org.jetbrains.kotlin.incremental.components.NoLookupLocation
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.TypeProjectionImpl
+import org.jetbrains.kotlin.types.TypeSubstitutor
+import org.jetbrains.kotlin.types.replace
 
 internal val CallableDescriptor.isSuspend: Boolean
     get() = this is FunctionDescriptor && isSuspend
@@ -65,3 +71,25 @@ fun CallableDescriptor.getOriginalParameter(parameter: ParameterDescriptor): Par
     this.extensionReceiverParameter -> this.original.extensionReceiverParameter!!
     else -> TODO("$parameter in $this")
 }
+
+fun KotlinType.replace(types: List<KotlinType>) = this.replace(types.map(::TypeProjectionImpl))
+
+fun FunctionDescriptor.substitute(vararg types: KotlinType): FunctionDescriptor {
+    val typeSubstitutor = TypeSubstitutor.create(
+            typeParameters
+                    .withIndex()
+                    .associateBy({ it.value.typeConstructor }, { TypeProjectionImpl(types[it.index]) })
+    )
+    return substitute(typeSubstitutor)!!
+}
+
+fun ClassDescriptor.getFunction(name: String, types: List<KotlinType>): FunctionDescriptor {
+    val typeSubstitutor = TypeSubstitutor.create(
+            declaredTypeParameters
+                    .withIndex()
+                    .associateBy({ it.value.typeConstructor }, { TypeProjectionImpl(types[it.index]) })
+    )
+    return unsubstitutedMemberScope
+            .getContributedFunctions(Name.identifier(name), NoLookupLocation.FROM_BACKEND).single().substitute(typeSubstitutor)!!
+}
+

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/KonanPlatform.kt

@@ -41,6 +41,10 @@ class KonanBuiltIns(storageManager: StorageManager) : KotlinBuiltIns(storageMana
     override fun getClassDescriptorFactories() =
             super.getClassDescriptorFactories() + KonanBuiltInClassDescriptorFactory(storageManager, builtInsModule)
 
+    override fun getSuspendFunction(parameterCount: Int): ClassDescriptor {
+        return getBuiltInClassByName(Name.identifier("SuspendFunction$parameterCount"))
+    }
+
     object FqNames {
         val packageName = FqName("konan.internal")
 

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/CallableReferenceLowering.kt

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.backend.common.descriptors.explicitParameters
 import org.jetbrains.kotlin.backend.common.lower.*
 import org.jetbrains.kotlin.backend.konan.Context
 import org.jetbrains.kotlin.backend.konan.descriptors.getKonanInternalClass
+import org.jetbrains.kotlin.backend.konan.descriptors.getKonanInternalFunctions
 import org.jetbrains.kotlin.backend.konan.descriptors.isFunctionOrKFunctionType
 import org.jetbrains.kotlin.backend.konan.descriptors.synthesizedName
 import org.jetbrains.kotlin.backend.konan.ir.createFakeOverrideDescriptor
@@ -47,16 +48,16 @@ import org.jetbrains.kotlin.ir.symbols.impl.IrSimpleFunctionSymbolImpl
 import org.jetbrains.kotlin.ir.util.createParameterDeclarations
 import org.jetbrains.kotlin.ir.util.getArguments
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
+import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.types.KotlinType
-import org.jetbrains.kotlin.types.TypeProjectionImpl
-import org.jetbrains.kotlin.types.TypeSubstitutor
-import org.jetbrains.kotlin.types.replace
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
 import org.jetbrains.kotlin.ir.util.transformFlat
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
+import org.jetbrains.kotlin.backend.common.descriptors.*
+import org.jetbrains.kotlin.types.*
 
 internal class CallableReferenceLowering(val context: Context): DeclarationContainerLoweringPass {
 
@@ -112,6 +113,17 @@ internal class CallableReferenceLowering(val context: Context): DeclarationConta
     private class BuiltFunctionReference(val functionReferenceClass: IrClass,
                                          val functionReferenceConstructor: IrConstructor)
 
+    private val COROUTINES_FQ_NAME = FqName.fromSegments(listOf("kotlin", "coroutines", "experimental"))
+    private val KOTLIN_FQ_NAME     = FqName("kotlin")
+
+    private val coroutinesScope    = context.irModule!!.descriptor.getPackage(COROUTINES_FQ_NAME).memberScope
+    private val kotlinPackageScope = context.irModule!!.descriptor.getPackage(KOTLIN_FQ_NAME).memberScope
+
+    private val continuationClassDescriptor = coroutinesScope
+            .getContributedClassifier(Name.identifier("Continuation"), NoLookupLocation.FROM_BACKEND) as ClassDescriptor
+
+    private val getContinuationDescriptor = context.builtIns.getKonanInternalFunctions("getContinuation").single()
+
     private inner class FunctionReferenceBuilder(val containingDeclaration: DeclarationDescriptor,
                                                  val functionReference: IrFunctionReference) {
 
@@ -125,21 +137,33 @@ internal class CallableReferenceLowering(val context: Context): DeclarationConta
         private lateinit var functionReferenceThis: IrValueParameterSymbol
         private lateinit var argumentToPropertiesMap: Map<ParameterDescriptor, IrFieldSymbol>
 
-        private fun KotlinType.replace(types: List<KotlinType>) = this.replace(types.map(::TypeProjectionImpl))
-
         private val kFunctionImplClassDescriptor = context.builtIns.getKonanInternalClass("KFunctionImpl")
 
         fun build(): BuiltFunctionReference {
             val startOffset = functionReference.startOffset
             val endOffset = functionReference.endOffset
 
+            val returnType = functionDescriptor.returnType!!
             val superTypes = mutableListOf(
-                    kFunctionImplClassDescriptor.defaultType.replace(listOf(functionDescriptor.returnType!!))
+                    kFunctionImplClassDescriptor.defaultType.replace(listOf(returnType))
             )
+
+            val numberOfParameters = unboundFunctionParameters.size
+            val functionClassDescriptor = context.reflectionTypes.getKFunction(numberOfParameters)
+            val functionParameterTypes = unboundFunctionParameters.map { it.type }
+            val functionClassTypeParameters = functionParameterTypes + returnType
+            superTypes += functionClassDescriptor.defaultType.replace(functionClassTypeParameters)
+
             var suspendFunctionClassDescriptor: ClassDescriptor? = null
-            val functionClassDescriptor = context.reflectionTypes.getKFunction(unboundFunctionParameters.size)
+            var suspendFunctionClassTypeParameters: List<KotlinType>? = null
-            val types = unboundFunctionParameters.map { it.type } + functionDescriptor.returnType!!
+            val lastParameterType = unboundFunctionParameters.lastOrNull()?.type
-            superTypes += functionClassDescriptor.defaultType.replace(types)
+            if (lastParameterType != null && TypeUtils.getClassDescriptor(lastParameterType) == continuationClassDescriptor) {
+                // If the last parameter is Continuation<> inherit from SuspendFunction.
+                suspendFunctionClassDescriptor = kotlinPackageScope.getContributedClassifier(
+                        Name.identifier("SuspendFunction${numberOfParameters - 1}"), NoLookupLocation.FROM_BACKEND) as ClassDescriptor
+                suspendFunctionClassTypeParameters = functionParameterTypes.dropLast(1) + lastParameterType.arguments.single().type
+                superTypes += suspendFunctionClassDescriptor.defaultType.replace(suspendFunctionClassTypeParameters)
+            }
 
             functionReferenceClassDescriptor = ClassDescriptorImpl(
                     /* containingDeclaration = */ containingDeclaration,
@@ -156,27 +180,29 @@ internal class CallableReferenceLowering(val context: Context): DeclarationConta
                     origin      = DECLARATION_ORIGIN_FUNCTION_REFERENCE_IMPL,
                     descriptor  = functionReferenceClassDescriptor
             )
+
             functionReferenceClass.createParameterDeclarations()
+
             functionReferenceThis = functionReferenceClass.thisReceiver!!.symbol
 
             val constructorBuilder = createConstructorBuilder()
 
-            val typeSubstitutor = TypeSubstitutor.create(
+            val invokeFunctionDescriptor = functionClassDescriptor.getFunction("invoke", functionClassTypeParameters)
-                    functionClassDescriptor.declaredTypeParameters
-                            .withIndex()
-                            .associateBy({ it.value.typeConstructor }, { TypeProjectionImpl(types[it.index]) })
-            )
-            val invokeFunctionDescriptor = functionClassDescriptor.unsubstitutedMemberScope
-                    .getContributedFunctions(Name.identifier("invoke"), NoLookupLocation.FROM_BACKEND).single().substitute(typeSubstitutor)!!
             val invokeMethodBuilder = createInvokeMethodBuilder(invokeFunctionDescriptor)
 
+            var suspendInvokeMethodBuilder: SymbolWithIrBuilder<IrSimpleFunctionSymbol, IrSimpleFunction>? = null
+            if (suspendFunctionClassDescriptor != null) {
+                val suspendInvokeFunctionDescriptor = suspendFunctionClassDescriptor.getFunction("invoke", suspendFunctionClassTypeParameters!!)
+                suspendInvokeMethodBuilder = createInvokeMethodBuilder(suspendInvokeFunctionDescriptor)
+            }
+
             val inheritedKFunctionImpl = kFunctionImplClassDescriptor.unsubstitutedMemberScope
                     .getContributedDescriptors()
                     .map { it.createFakeOverrideDescriptor(functionReferenceClassDescriptor) }
                     .filterNotNull()
             val contributedDescriptors = (
-                    inheritedKFunctionImpl + invokeMethodBuilder.symbol.descriptor
+                    inheritedKFunctionImpl + invokeMethodBuilder.symbol.descriptor + suspendInvokeMethodBuilder?.symbol?.descriptor
-                    ).toList()
+                    ).filterNotNull().toList()
             functionReferenceClassDescriptor.initialize(
                     SimpleMemberScope(contributedDescriptors), setOf(constructorBuilder.symbol.descriptor), null)
 
@@ -186,6 +212,11 @@ internal class CallableReferenceLowering(val context: Context): DeclarationConta
             invokeMethodBuilder.initialize()
             functionReferenceClass.declarations.add(invokeMethodBuilder.ir)
 
+            suspendInvokeMethodBuilder?.let {
+                it.initialize()
+                functionReferenceClass.declarations.add(it.ir)
+            }
+
             return BuiltFunctionReference(functionReferenceClass, constructorBuilder.ir)
         }
 
@@ -266,7 +297,7 @@ internal class CallableReferenceLowering(val context: Context): DeclarationConta
                 }
             }
 
-        private fun createInvokeMethodBuilder(functionInvokeFunctionDescriptor: FunctionDescriptor)
+        private fun createInvokeMethodBuilder(superFunctionDescriptor: FunctionDescriptor)
                 = object : SymbolWithIrBuilder<IrSimpleFunctionSymbol, IrSimpleFunction>() {
 
             override fun buildSymbol() = IrSimpleFunctionSymbolImpl(
@@ -281,7 +312,7 @@ internal class CallableReferenceLowering(val context: Context): DeclarationConta
 
             override fun doInitialize() {
                 val descriptor = symbol.descriptor as SimpleFunctionDescriptorImpl
-                val valueParameters = functionInvokeFunctionDescriptor.valueParameters
+                val valueParameters = superFunctionDescriptor.valueParameters
                         .map { it.copyAsValueParameter(descriptor, it.index) }
 
                 descriptor.initialize(
@@ -289,21 +320,23 @@ internal class CallableReferenceLowering(val context: Context): DeclarationConta
                         /* dispatchReceiverParameter    = */ functionReferenceClassDescriptor.thisAsReceiverParameter,
                         /* typeParameters               = */ emptyList(),
                         /* unsubstitutedValueParameters = */ valueParameters,
-                        /* unsubstitutedReturnType      = */ functionInvokeFunctionDescriptor.returnType,
+                        /* unsubstitutedReturnType      = */ superFunctionDescriptor.returnType,
                         /* modality                     = */ Modality.FINAL,
                         /* visibility                   = */ Visibilities.PRIVATE).apply {
-                    overriddenDescriptors += functionInvokeFunctionDescriptor
+                    overriddenDescriptors              +=    superFunctionDescriptor
+                    isSuspend                           =    superFunctionDescriptor.isSuspend
                 }
             }
 
             override fun buildIr(): IrSimpleFunction {
                 val startOffset = functionReference.startOffset
                 val endOffset = functionReference.endOffset
+                val ourSymbol = symbol
                 return IrFunctionImpl(
                         startOffset = startOffset,
                         endOffset   = endOffset,
                         origin      = DECLARATION_ORIGIN_FUNCTION_REFERENCE_IMPL,
-                        symbol      = symbol).apply {
+                        symbol      = ourSymbol).apply {
 
                     val function = this
                     val irBuilder = context.createIrBuilder(function.symbol, startOffset, endOffset)
@@ -316,10 +349,17 @@ internal class CallableReferenceLowering(val context: Context): DeclarationConta
                                     var unboundIndex = 0
                                     val unboundArgsSet = unboundFunctionParameters.toSet()
                                     functionParameters.forEach {
-                                        val argument = if (unboundArgsSet.contains(it))
+                                        val argument =
-                                            irGet(valueParameters[unboundIndex++].symbol)
+                                                if (!unboundArgsSet.contains(it))
-                                        else
+                                                    // Bound parameter - read from field.
-                                            irGetField(irGet(functionReferenceThis), argumentToPropertiesMap[it]!!)
+                                                    irGetField(irGet(functionReferenceThis), argumentToPropertiesMap[it]!!)
+                                                else {
+                                                    if (ourSymbol.descriptor.isSuspend && unboundIndex == valueParameters.size)
+                                                        // For suspend functions the last argument is continuation and it is implicit.
+                                                        irCall(getContinuationDescriptor.substitute(ourSymbol.descriptor.returnType!!))
+                                                    else
+                                                        irGet(valueParameters[unboundIndex++].symbol)
+                                                }
                                         when (it) {
                                             functionDescriptor.dispatchReceiverParameter -> dispatchReceiver = argument
                                             functionDescriptor.extensionReceiverParameter -> extensionReceiver = argument

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/SuspendFunctionsLowering.kt

@@ -50,11 +50,9 @@ import org.jetbrains.kotlin.ir.visitors.*
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.types.KotlinType
-import org.jetbrains.kotlin.types.TypeProjectionImpl
-import org.jetbrains.kotlin.types.TypeSubstitutor
-import org.jetbrains.kotlin.types.replace
 import org.jetbrains.kotlin.types.typeUtil.isUnit
 import org.jetbrains.kotlin.types.typeUtil.makeNotNullable
+import org.jetbrains.kotlin.backend.common.descriptors.*
 
 internal class SuspendFunctionsLowering(val context: Context): DeclarationContainerLoweringPass {
 
@@ -236,17 +234,6 @@ internal class SuspendFunctionsLowering(val context: Context): DeclarationContai
     private val getContinuationDescriptor = context.builtIns.getKonanInternalFunctions("getContinuation").single()
     private val returnIfSuspendedDescriptor = context.builtIns.getKonanInternalFunctions("returnIfSuspended").single()
 
-    private fun KotlinType.replace(types: List<KotlinType>) = this.replace(types.map(::TypeProjectionImpl))
-
-    private fun FunctionDescriptor.substitute(vararg types: KotlinType): FunctionDescriptor {
-        val typeSubstitutor = TypeSubstitutor.create(
-                typeParameters
-                        .withIndex()
-                        .associateBy({ it.value.typeConstructor }, { TypeProjectionImpl(types[it.index]) })
-        )
-        return substitute(typeSubstitutor)!!
-    }
-
     private fun removeReturnIfSuspendedCall(irFunction: IrFunction) {
         irFunction.transformChildrenVoid(object: IrElementTransformerVoid() {
             override fun visitCall(expression: IrCall): IrExpression {
@@ -332,22 +319,23 @@ internal class SuspendFunctionsLowering(val context: Context): DeclarationContai
             val superTypes = mutableListOf<KotlinType>(coroutineImplClassDescriptor.defaultType)
             var suspendFunctionClassDescriptor: ClassDescriptor? = null
             var functionClassDescriptor: ClassDescriptor? = null
+            var suspendFunctionClassTypeArguments: List<KotlinType>? = null
+            var functionClassTypeArguments: List<KotlinType>? = null
             if (unboundFunctionParameters != null) {
                 // Suspend lambda inherits SuspendFunction.
                 val numberOfParameters = unboundFunctionParameters.size
                 suspendFunctionClassDescriptor = kotlinPackageScope.getContributedClassifier(
                         Name.identifier("SuspendFunction$numberOfParameters"), NoLookupLocation.FROM_BACKEND) as ClassDescriptor
-                superTypes += suspendFunctionClassDescriptor.defaultType.replace(
+                val unboundParameterTypes = unboundFunctionParameters.map { it.type }
-                        unboundFunctionParameters.map { it.type } + irFunction.descriptor.returnType!!
+                suspendFunctionClassTypeArguments = unboundParameterTypes + irFunction.descriptor.returnType!!
-                )
+                superTypes += suspendFunctionClassDescriptor.defaultType.replace(suspendFunctionClassTypeArguments)
 
                 functionClassDescriptor = kotlinPackageScope.getContributedClassifier(
                         Name.identifier("Function${numberOfParameters + 1}"), NoLookupLocation.FROM_BACKEND) as ClassDescriptor
                 val continuationType = continuationClassDescriptor.defaultType.replace(listOf(irFunction.descriptor.returnType!!))
-                superTypes += functionClassDescriptor.defaultType.replace(
+                functionClassTypeArguments = unboundParameterTypes + continuationType + context.builtIns.nullableAnyType
-                        unboundFunctionParameters.map { it.type }     // Main arguments,
+                superTypes += functionClassDescriptor.defaultType.replace(functionClassTypeArguments)
-                                + continuationType                    // and continuation.
+
-                                + context.builtIns.nullableAnyType)   // Return type
             }
             coroutineClassDescriptor = ClassDescriptorImpl(
                     /* containingDeclaration = */ irFunction.descriptor.containingDeclaration,
@@ -397,10 +385,8 @@ internal class SuspendFunctionsLowering(val context: Context): DeclarationContai
                 if (createFunctionDescriptor != null)
                     overriddenMap += createFunctionDescriptor to createMethodBuilder.symbol.descriptor
 
-                val invokeFunctionDescriptor = functionClassDescriptor!!.unsubstitutedMemberScope
+                val invokeFunctionDescriptor = functionClassDescriptor!!.getFunction("invoke", functionClassTypeArguments!!)
-                        .getContributedFunctions(Name.identifier("invoke"), NoLookupLocation.FROM_BACKEND).single()
+                val suspendInvokeFunctionDescriptor = suspendFunctionClassDescriptor!!.getFunction("invoke", suspendFunctionClassTypeArguments!!)
-                val suspendInvokeFunctionDescriptor = suspendFunctionClassDescriptor!!.unsubstitutedMemberScope
-                        .getContributedFunctions(Name.identifier("invoke"), NoLookupLocation.FROM_BACKEND).single()
                 invokeMethodBuilder = createInvokeMethodBuilder(
                         suspendFunctionInvokeFunctionDescriptor = suspendInvokeFunctionDescriptor,
                         functionInvokeFunctionDescriptor        = invokeFunctionDescriptor,