[Native] Support suspend functions as supertypes

^KT-46777

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/AbstractSuspendFunctionsLowering.kt

@@ -20,14 +20,12 @@ import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.IrInstanceInitializerCallImpl
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
-import org.jetbrains.kotlin.ir.types.IrType
-import org.jetbrains.kotlin.ir.types.classifierOrFail
-import org.jetbrains.kotlin.ir.types.defaultType
+import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl
-import org.jetbrains.kotlin.ir.types.typeWith
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.*
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.util.OperatorNameConventions
 
 abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val context: C) : FileLoweringPass {
 
@@ -58,6 +56,62 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
         markSuspendLambdas(irFile)
         buildCoroutines(irFile)
         transformCallableReferencesToSuspendLambdas(irFile)
+        addMissingSupertypesToSuspendFunctionImplementingClasses(irFile)
+    }
+
+    private fun addMissingSupertypesToSuspendFunctionImplementingClasses(irFile: IrFile) {
+        irFile.acceptChildrenVoid(object : IrElementVisitorVoid {
+            override fun visitElement(element: IrElement) {
+                // Don't need to iterate through children. All local classes are already moved to the top level by this moment.
+            }
+
+            override fun visitClass(declaration: IrClass) {
+                addMissingSupertypes(declaration)
+                declaration.acceptChildrenVoid(this)
+            }
+
+            private fun addMissingSupertypes(clazz: IrClass) {
+                val suspendFunctionTypes = mutableSetOf<IrSimpleType>()
+                for (superType in getAllSubstitutedSupertypes(clazz)) {
+                    when {
+                        superType.isFunctionMarker() -> Unit // Proceed with others.
+                        superType.isFunction() -> {
+                            // Mixing suspend and non-suspend function supertypes is not allowed by the frontend. So can stop here.
+                            return
+                        }
+                        superType.isSuspendFunction() -> suspendFunctionTypes += superType
+                    }
+                }
+
+                for (suspendFunctionType in suspendFunctionTypes) {
+                    val suspendFunctionClassSymbol = suspendFunctionType.classOrNull ?: continue
+                    val suspendFunctionSymbol = suspendFunctionClassSymbol.owner.simpleFunctions().single {
+                        it.name == OperatorNameConventions.INVOKE
+                    }.symbol
+
+                    val invokeFunction = clazz.simpleFunctions().single {
+                        it.name == OperatorNameConventions.INVOKE && suspendFunctionSymbol in it.overriddenSymbols
+                    }
+
+                    val suspendFunctionArity = suspendFunctionSymbol.owner.valueParameters.size
+                    val functionClassSymbol = symbols.functionN(suspendFunctionArity + 1)
+                    val functionSymbol = functionClassSymbol.owner.simpleFunctions().single {
+                        it.name == OperatorNameConventions.INVOKE
+                    }.symbol
+
+                    invokeFunction.overriddenSymbols += functionSymbol
+
+                    val functionClassTypeArguments = suspendFunctionType.arguments.mapIndexed { index, argument ->
+                        val type = (argument as IrTypeProjection).type
+                        if (index == suspendFunctionArity) continuationClassSymbol.typeWith(type) else type
+                    } + context.irBuiltIns.anyNType
+
+                    val functionType = functionClassSymbol.typeWith(functionClassTypeArguments)
+
+                    clazz.superTypes += functionType
+                }
+            }
+        })
     }
 
     private fun buildCoroutines(irFile: IrFile) {
@@ -96,7 +150,8 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
                 expression.acceptChildrenVoid(this)
 
                 if (expression.isSuspend) {
-                    suspendLambdas[expression.symbol.owner] = expression
+                    val old = suspendLambdas.put(expression.symbol.owner, expression)
+                    if (old != null) error("Rewriting $old")
                 }
             }
         })
@@ -228,7 +283,8 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
 
     private fun buildCoroutine(irFunction: IrSimpleFunction, functionReference: IrFunctionReference?): IrClass {
         val coroutine = CoroutineBuilder(irFunction, functionReference).build()
-        builtCoroutines[irFunction] = coroutine
+        val old = builtCoroutines.put(irFunction, coroutine)
+        if (old != null) error("Rewriting $old")
 
         if (functionReference == null) {
             // It is not a lambda - replace original function with a call to constructor of the built coroutine.
@@ -294,7 +350,7 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
         private val coroutineBaseClass = getCoroutineBaseClass(irFunction)
         private val coroutineBaseClassConstructor = coroutineBaseClass.owner.constructors.single { it.valueParameters.size == 1 }
         private val create1Function = coroutineBaseClass.owner.simpleFunctions()
-            .single { it.name.asString() == "create" && it.valueParameters.size == 1 }
+            .single { it.name == CREATE_IDENTIFIER && it.valueParameters.size == 1 }
         private val create1CompletionParameter = create1Function.valueParameters[0]
 
         private val coroutineConstructors = mutableListOf<IrConstructor>()
@@ -330,7 +386,7 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
                 coroutineFactoryConstructor = buildFactoryConstructor(boundFunctionParameters!!)
 
                 val createFunctionSymbol = coroutineBaseClass.owner.simpleFunctions()
-                    .atMostOne { it.name.asString() == "create" && it.valueParameters.size == unboundFunctionParameters!!.size + 1 }
+                    .atMostOne { it.name == CREATE_IDENTIFIER && it.valueParameters.size == unboundFunctionParameters!!.size + 1 }
                     ?.symbol
 
                 createMethod = buildCreateMethod(
@@ -340,9 +396,9 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
                 )
 
                 val invokeFunctionSymbol =
-                    functionClass!!.simpleFunctions().single { it.name.asString() == "invoke" }.symbol
+                    functionClass!!.simpleFunctions().single { it.name == OperatorNameConventions.INVOKE }.symbol
                 val suspendInvokeFunctionSymbol =
-                    suspendFunctionClass!!.simpleFunctions().single { it.name.asString() == "invoke" }.symbol
+                    suspendFunctionClass!!.simpleFunctions().single { it.name == OperatorNameConventions.INVOKE }.symbol
 
                 buildInvokeMethod(
                     suspendFunctionInvokeFunctionSymbol = suspendInvokeFunctionSymbol,
@@ -447,7 +503,7 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
             startOffset = irFunction.startOffset
             endOffset = irFunction.endOffset
             origin = DECLARATION_ORIGIN_COROUTINE_IMPL
-            name = Name.identifier("create")
+            name = CREATE_IDENTIFIER
             visibility = DescriptorVisibilities.PROTECTED
             returnType = coroutineClass.defaultType
         }.apply {
@@ -503,7 +559,7 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
             startOffset = irFunction.startOffset
             endOffset = irFunction.endOffset
             origin = DECLARATION_ORIGIN_COROUTINE_IMPL
-            name = Name.identifier("invoke")
+            name = OperatorNameConventions.INVOKE
             visibility = DescriptorVisibilities.PROTECTED
             returnType = context.irBuiltIns.anyNType
             isSuspend = true
@@ -629,4 +685,8 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
             addChild(it)
         }
     }
+
+    companion object {
+        private val CREATE_IDENTIFIER = Name.identifier("create")
+    }
 }

compiler/ir/backend.common/src/org/jetbrains/kotlin/ir/util/IrTypeUtils.kt

@@ -26,22 +26,26 @@ val kotlinPackageFqn = FqName.fromSegments(listOf("kotlin"))
 private val kotlinReflectionPackageFqn = kotlinPackageFqn.child(Name.identifier("reflect"))
 private val kotlinCoroutinesPackageFqn = kotlinPackageFqn.child(Name.identifier("coroutines"))
 
+fun IrType.isFunctionMarker(): Boolean = classifierOrNull?.isClassWithName("Function", kotlinPackageFqn) == true
 fun IrType.isFunction(): Boolean = classifierOrNull?.isClassWithNamePrefix("Function", kotlinPackageFqn) == true
 fun IrType.isKFunction(): Boolean = classifierOrNull?.isClassWithNamePrefix("KFunction", kotlinReflectionPackageFqn) == true
 fun IrType.isSuspendFunction(): Boolean = classifierOrNull?.isClassWithNamePrefix("SuspendFunction", kotlinCoroutinesPackageFqn) == true
 fun IrType.isKSuspendFunction(): Boolean = classifierOrNull?.isClassWithNamePrefix("KSuspendFunction", kotlinReflectionPackageFqn) == true
 
+fun IrClassifierSymbol.isFunctionMarker(): Boolean = this.isClassWithName("Function", kotlinPackageFqn)
 fun IrClassifierSymbol.isFunction(): Boolean = this.isClassWithNamePrefix("Function", kotlinPackageFqn)
 fun IrClassifierSymbol.isKFunction(): Boolean = this.isClassWithNamePrefix("KFunction", kotlinReflectionPackageFqn)
+fun IrClassifierSymbol.isSuspendFunction(): Boolean = this.isClassWithNamePrefix("SuspendFunction", kotlinCoroutinesPackageFqn)
 fun IrClassifierSymbol.isKSuspendFunction(): Boolean = this.isClassWithNamePrefix("KSuspendFunction", kotlinReflectionPackageFqn)
 
+private fun IrClassifierSymbol.isClassWithName(name: String, packageFqName: FqName): Boolean {
+    val declaration = owner as IrDeclarationWithName
+    return name == declaration.name.asString() && (declaration.parent as? IrPackageFragment)?.fqName == packageFqName
+}
+
 private fun IrClassifierSymbol.isClassWithNamePrefix(prefix: String, packageFqName: FqName): Boolean {
     val declaration = owner as IrDeclarationWithName
-    val name = declaration.name.asString()
-    if (!name.startsWith(prefix)) return false
-    val parent = declaration.parent as? IrPackageFragment ?: return false
-
-    return parent.fqName == packageFqName
+    return declaration.name.asString().startsWith(prefix) && (declaration.parent as? IrPackageFragment)?.fqName == packageFqName
 }
 
 fun IrType.superTypes(): List<IrType> = classifierOrNull?.superTypes() ?: emptyList()

compiler/testData/codegen/box/coroutines/suspendFunctionAsSupertype.kt

@@ -4,28 +4,170 @@
 
 import kotlin.coroutines.*
 
-var res = "FAIL"
+var failure: String? = "FAIL ILLEGAL STATE"
 
-class C: suspend () -> Unit {
-    override suspend fun invoke() {
-        res = "OK"
+class TopLevel1: suspend () -> Int {
+    override suspend fun invoke(): Int {
+        failure = null
+        return 42
     }
 }
 
-fun box(): String {
-    val o: suspend () -> Unit = object : suspend () -> Unit {
-        override suspend fun invoke() {
-            res = "OK"
+class TopLevel2: suspend (String) -> Int {
+    override suspend fun invoke(p: String): Int {
+        failure = null
+        return p.length
+    }
+}
+
+class Outer {
+    class Nested1: suspend () -> Int {
+        override suspend fun invoke(): Int {
+            failure = null
+            return 42
         }
     }
-    o.startCoroutine(Continuation(EmptyCoroutineContext) {
-        it.getOrThrow()
-    })
-    if (res != "OK") return "FAIL 1: $res"
 
-    res = "FAIL 2"
-    C().startCoroutine(Continuation(EmptyCoroutineContext) {
-        it.getOrThrow()
-    })
-    return res
-}
+    class Nested2: suspend (String) -> Int {
+        override suspend fun invoke(p: String): Int {
+            failure = null
+            return p.length
+        }
+    }
+}
+
+interface I : suspend (String) -> Int
+abstract class Abstract1 : suspend (String) -> Int
+abstract class Abstract2 : suspend (String) -> Int {
+    override suspend fun invoke(p: String): Int {
+        failure = null
+        return p.length
+    }
+}
+
+fun objectLiteral1(): String? {
+    failure = "FAIL OBJECT LITERAL 1"
+    val o: suspend () -> Int = object : suspend () -> Int {
+        override suspend fun invoke(): Int {
+            failure = null
+            return 42
+        }
+    }
+    o.startCoroutine(Continuation(EmptyCoroutineContext) { it.getOrThrow() })
+    return failure
+}
+
+fun objectLiteral2(): String? {
+    failure = "FAIL OBJECT LITERAL 2"
+    val o: suspend (String) -> Int = object : suspend (String) -> Int {
+        override suspend fun invoke(p: String): Int {
+            failure = null
+            return p.length
+        }
+    }
+    o.startCoroutine("Hello", Continuation(EmptyCoroutineContext) { it.getOrThrow() })
+    return failure
+}
+
+fun topLevelClass1(): String? {
+    failure = "FAIL TOP LEVEL CLASS 1"
+    TopLevel1().startCoroutine(Continuation(EmptyCoroutineContext) { it.getOrThrow() })
+    return failure
+}
+
+fun topLevelClass2(): String? {
+    failure = "FAIL TOP LEVEL CLASS 2"
+    TopLevel2().startCoroutine("Hello", Continuation(EmptyCoroutineContext) { it.getOrThrow() })
+    return failure
+}
+
+fun nestedClass1(): String? {
+    failure = "FAIL NESTED CLASS 1"
+    Outer.Nested1().startCoroutine(Continuation(EmptyCoroutineContext) { it.getOrThrow() })
+    return failure
+}
+
+fun nestedClass2(): String? {
+    failure = "FAIL NESTED CLASS 2"
+    Outer.Nested2().startCoroutine("Hello", Continuation(EmptyCoroutineContext) { it.getOrThrow() })
+    return failure
+}
+
+fun localClass1(): String? {
+    failure = "FAIL LOCAL CLASS 1"
+    class Local : suspend () -> Int {
+        override suspend fun invoke(): Int {
+            failure = null
+            return 42
+        }
+    }
+    Local().startCoroutine(Continuation(EmptyCoroutineContext) { it.getOrThrow() })
+    return failure
+}
+
+fun localClass2(): String? {
+    fun foo(): String? {
+        fun bar(): String? {
+            failure = "FAIL LOCAL CLASS 2"
+            class Local : suspend (String) -> Int {
+                override suspend fun invoke(p: String): Int {
+                    failure = null
+                    return p.length
+                }
+            }
+            Local().startCoroutine("Hello", Continuation(EmptyCoroutineContext) { it.getOrThrow() })
+            return failure
+        }
+        return bar()
+    }
+    return foo()
+}
+
+fun inherited1(): String? {
+    failure = "FAIL INHERITED 1"
+    val o = object : I {
+        override suspend fun invoke(p: String): Int {
+            failure = null
+            return p.length
+        }
+    }
+    o.startCoroutine("Hello", Continuation(EmptyCoroutineContext) { it.getOrThrow() })
+    return failure
+}
+
+fun inherited2(): String? {
+    failure = "FAIL INHERITED 2"
+    val o = object : I, Abstract1() {
+        override suspend fun invoke(p: String): Int {
+            failure = null
+            return p.length
+        }
+    }
+    o.startCoroutine("Hello", Continuation(EmptyCoroutineContext) { it.getOrThrow() })
+    return failure
+}
+
+fun inherited3(): String? {
+    failure = "FAIL INHERITED 3"
+    val o = object : I, Abstract2() {}
+    o.startCoroutine("Hello", Continuation(EmptyCoroutineContext) { it.getOrThrow() })
+    return failure
+}
+
+fun box(): String {
+    val failures = listOfNotNull(
+        objectLiteral1(),
+        objectLiteral2(),
+        topLevelClass1(),
+        topLevelClass2(),
+        nestedClass1(),
+        nestedClass2(),
+        localClass1(),
+        localClass2(),
+        inherited1(),
+        inherited2(),
+        inherited3()
+    )
+
+    return if (failures.isNotEmpty()) failures.joinToString("\n") else "OK"
+}

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

@@ -16,11 +16,12 @@ import org.jetbrains.kotlin.backend.common.push
 import org.jetbrains.kotlin.backend.konan.*
 import org.jetbrains.kotlin.backend.konan.cgen.*
 import org.jetbrains.kotlin.backend.konan.descriptors.allOverriddenFunctions
+import org.jetbrains.kotlin.backend.konan.ir.*
 import org.jetbrains.kotlin.backend.konan.ir.companionObject
 import org.jetbrains.kotlin.backend.konan.ir.getSuperClassNotAny
-import org.jetbrains.kotlin.backend.konan.ir.isReal
 import org.jetbrains.kotlin.backend.konan.llvm.IntrinsicType
 import org.jetbrains.kotlin.backend.konan.llvm.tryGetIntrinsicType
+import org.jetbrains.kotlin.builtins.functions.FunctionClassKind
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
@@ -28,6 +29,7 @@ import org.jetbrains.kotlin.descriptors.isFinalClass
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.symbols.IrTypeParameterSymbol
 import org.jetbrains.kotlin.ir.types.*
@@ -75,10 +77,6 @@ private class BackendChecker(val context: Context, val irFile: IrFile) : IrEleme
     }
 
     override fun visitDeclaration(declaration: IrDeclarationBase) {
-        if (declaration is IrClass && declaration.isKotlinObjCClass()) {
-            checkKotlinObjCClass(declaration)
-        }
-
         outerDeclarations.push(declaration)
         try {
             super.visitDeclaration(declaration)
@@ -588,6 +586,46 @@ private class BackendChecker(val context: Context, val irFile: IrFile) : IrEleme
         typeParameter.superTypes.forEach { checkIrKType(irElement, it, seenTypeParameters) }
         seenTypeParameters.remove(typeParameter)
     }
+
+    override fun visitClass(declaration: IrClass) {
+        super.visitClass(declaration)
+
+        if (declaration.isKotlinObjCClass()) {
+            checkKotlinObjCClass(declaration)
+        }
+        checkNotImplementsMultipleSuspendLambdas(declaration)
+    }
+}
+
+// TODO: remove this check in KT-46775
+private fun BackendChecker.checkNotImplementsMultipleSuspendLambdas(clazz: IrClass) {
+    var suspendFunctionClassSymbol: IrClassSymbol? = null // Can be null (no suspend function supertypes) or non-null (the single one).
+
+    fun checkSupertype(superType: IrType): Boolean {
+        val superTypeClassSymbol = superType.classOrNull
+                ?: return true // This is not a class. Continue with others.
+
+        return when {
+            superTypeClassSymbol.isFunctionMarker() -> true // Continue with others.
+            superTypeClassSymbol.isFunction() -> {
+                // Mixing suspend and non-suspend function supertypes is not allowed by the frontend. So can stop here.
+                false
+            }
+            superTypeClassSymbol.isSuspendFunction() -> {
+                suspendFunctionClassSymbol = when (suspendFunctionClassSymbol) {
+                    null, superTypeClassSymbol -> superTypeClassSymbol
+                    else -> reportError(clazz, "Multiple suspend function types as supertypes are not supported yet")
+                }
+                true // Continue with others.
+            }
+            else -> {
+                // Recurse with supertypes.
+                superTypeClassSymbol.owner.superTypes.all(::checkSupertype)
+            }
+        }
+    }
+
+    clazz.superTypes.all(::checkSupertype)
 }
 
 private fun BackendChecker.checkCanGenerateCCall(expression: IrCall, isInvoke: Boolean) {