[K/N] Remove hack from suspend function inheritance to allow export

Before that commit, function in coroutine implementation class was
override for both SuspendFunctionN's and FunctionN+1's invoke,
which somehow works, but breaks exporting classes with such functions.

Now it's two separate functions, as it should be normally.

^KT-49395

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

@@ -6,15 +6,13 @@
 package org.jetbrains.kotlin.backend.common.lower
 
 import org.jetbrains.kotlin.backend.common.*
+import org.jetbrains.kotlin.backend.common.descriptors.synthesizedName
 import org.jetbrains.kotlin.backend.common.ir.*
 import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.builders.*
-import org.jetbrains.kotlin.ir.builders.declarations.buildClass
+import org.jetbrains.kotlin.ir.builders.declarations.*
-import org.jetbrains.kotlin.ir.builders.declarations.buildConstructor
-import org.jetbrains.kotlin.ir.builders.declarations.buildField
-import org.jetbrains.kotlin.ir.builders.declarations.buildFun
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.IrInstanceInitializerCallImpl
@@ -35,6 +33,9 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
     protected abstract val stateMachineMethodName: Name
     protected abstract fun getCoroutineBaseClass(function: IrFunction): IrClassSymbol
     protected abstract fun nameForCoroutineClass(function: IrFunction): Name
+    protected abstract fun IrBuilderWithScope.launchSuspendFunctionWithGivenContinuation(
+        symbol: IrSimpleFunctionSymbol, dispatchReceiver: IrExpression,
+        arguments: List<IrExpression>, continuation: IrExpression) : IrExpression
 
     protected abstract fun buildStateMachine(
         stateMachineFunction: IrFunction,
@@ -85,22 +86,24 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
 
                 for (suspendFunctionType in suspendFunctionTypes) {
                     val suspendFunctionClassSymbol = suspendFunctionType.classOrNull ?: continue
-                    val suspendFunctionSymbol = suspendFunctionClassSymbol.owner.simpleFunctions().single {
+                    val suspendFunction = 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 invokeFunction = clazz.simpleFunctions().single {
+                        it.name == OperatorNameConventions.INVOKE && it.overrides(suspendFunction)
+                    }
+
+                    if (invokeFunction.modality == Modality.ABSTRACT) {
+                        continue
+                    }
+
+                    val suspendFunctionArity = suspendFunction.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
@@ -109,6 +112,51 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
                     val functionType = functionClassSymbol.typeWith(functionClassTypeArguments)
 
                     clazz.superTypes += functionType
+
+                    context.irFactory.buildFun {
+                        startOffset = invokeFunction.startOffset
+                        endOffset = invokeFunction.endOffset
+                        origin = DECLARATION_ORIGIN_COROUTINE_IMPL
+                        name = OperatorNameConventions.INVOKE
+                        visibility = DescriptorVisibilities.PROTECTED
+                        returnType = context.irBuiltIns.anyNType
+                    }.apply {
+                        parent = clazz
+                        clazz.declarations += this
+
+                        typeParameters = invokeFunction.typeParameters.map { parameter ->
+                            parameter.copyToWithoutSuperTypes(this, origin = DECLARATION_ORIGIN_COROUTINE_IMPL)
+                                .apply { superTypes += parameter.superTypes }
+                        }
+
+                        valueParameters = invokeFunction.valueParameters
+                            .mapIndexed { index, parameter ->
+                                parameter.copyTo(this, DECLARATION_ORIGIN_COROUTINE_IMPL, index)
+                            }
+                        valueParameters += buildValueParameter(this) {
+                            index = valueParameters.size
+                            type = context.irBuiltIns.anyNType
+                            name = "completion".synthesizedName
+                        }
+
+                        this.createDispatchReceiverParameter()
+
+                        overriddenSymbols += functionSymbol
+
+                        val thisReceiver = dispatchReceiverParameter!!
+
+                        val irBuilder = context.createIrBuilder(symbol, startOffset, endOffset)
+                        body = irBuilder.irBlockBody(startOffset, endOffset) {
+                            +irReturn(
+                                launchSuspendFunctionWithGivenContinuation(
+                                    invokeFunction.symbol,
+                                    irGet(thisReceiver),
+                                    valueParameters.dropLast(1).map { irGet(it) },
+                                    irGet(valueParameters.last())
+                                )
+                            )
+                        }
+                    }
                 }
             }
         })
@@ -356,9 +404,7 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
         fun build(): BuiltCoroutine {
             val superTypes = mutableListOf(coroutineBaseClass.defaultType)
             var suspendFunctionClass: IrClass? = null
-            var functionClass: IrClass? = null
             val suspendFunctionClassTypeArguments: List<IrType>?
-            val functionClassTypeArguments: List<IrType>?
             if (unboundFunctionParameters != null) {
                 // Suspend lambda inherits SuspendFunction.
                 val numberOfParameters = unboundFunctionParameters.size
@@ -366,10 +412,6 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
                 val unboundParameterTypes = unboundFunctionParameters.map { it.type }
                 suspendFunctionClassTypeArguments = unboundParameterTypes + irFunction.returnType
                 superTypes += suspendFunctionClass.typeWith(suspendFunctionClassTypeArguments)
-
-                functionClass = symbols.functionN(numberOfParameters + 1).owner
-                functionClassTypeArguments = unboundParameterTypes + continuationType + context.irBuiltIns.anyNType
-                superTypes += functionClass.typeWith(functionClassTypeArguments)
             }
 
             val coroutineConstructor = buildConstructor()
@@ -393,16 +435,13 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
                     coroutineConstructor = coroutineConstructor
                 )
 
-                val invokeFunctionSymbol =
+                val suspendInvokeFunction =
-                    functionClass!!.simpleFunctions().single { it.name == OperatorNameConventions.INVOKE }.symbol
+                    suspendFunctionClass!!.simpleFunctions().single { it.name == OperatorNameConventions.INVOKE }
-                val suspendInvokeFunctionSymbol =
-                    suspendFunctionClass!!.simpleFunctions().single { it.name == OperatorNameConventions.INVOKE }.symbol
 
                 buildInvokeMethod(
-                    suspendFunctionInvokeFunctionSymbol = suspendInvokeFunctionSymbol,
+                    functionInvokeFunction = suspendInvokeFunction,
-                    functionInvokeFunctionSymbol = invokeFunctionSymbol,
                     createFunction = createMethod,
-                    stateMachineFunction = invokeSuspendMethod
+                    stateMachineFunction = invokeSuspendMethod,
                 )
             }
 
@@ -549,8 +588,7 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
         }
 
         private fun buildInvokeMethod(
-            suspendFunctionInvokeFunctionSymbol: IrSimpleFunctionSymbol,
+            functionInvokeFunction: IrSimpleFunction,
-            functionInvokeFunctionSymbol: IrSimpleFunctionSymbol,
             createFunction: IrFunction,
             stateMachineFunction: IrFunction
         ): IrSimpleFunction = context.irFactory.buildFun {
@@ -579,8 +617,7 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
 
             this.createDispatchReceiverParameter()
 
-            overriddenSymbols += functionInvokeFunctionSymbol
+            overriddenSymbols += functionInvokeFunction.symbol
-            overriddenSymbols += suspendFunctionInvokeFunctionSymbol
 
             val thisReceiver = dispatchReceiverParameter!!
 

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

@@ -251,12 +251,10 @@ internal class FunctionReferenceLowering(val context: Context): FileLoweringPass
             } else {
                 val numberOfParameters = unboundFunctionParameters.size
                 val functionParameterTypes = unboundFunctionParameters.map { it.type }
-                val functionClass: IrClass
+                val functionClass: IrClass?
                 val suspendFunctionClass: IrClass?
                 if (isKSuspendFunction) {
-                    functionClass = symbols.functionN(numberOfParameters + 1).owner
+                    functionClass = null
-                    val continuationType = continuationClassSymbol.typeWith(referencedFunction.returnType)
-                    superTypes += functionClass.typeWith(functionParameterTypes + continuationType + irBuiltIns.anyNType)
                     suspendFunctionClass = symbols.kSuspendFunctionN(numberOfParameters).owner
                     superTypes += suspendFunctionClass.typeWith(functionParameterTypes + referencedFunction.returnType)
                 } else {
@@ -275,13 +273,10 @@ internal class FunctionReferenceLowering(val context: Context): FileLoweringPass
                     }
                 }
 
-                if (!isKSuspendFunction)
+                if (functionClass != null)
                     buildInvokeMethod(functionClass.getInvokeFunction())
                 if (suspendFunctionClass != null) {
-                    buildInvokeMethod(suspendFunctionClass.getInvokeFunction()).also {
+                    buildInvokeMethod(suspendFunctionClass.getInvokeFunction())
-                        if (isKSuspendFunction)
-                            it.overriddenSymbols += functionClass.getInvokeFunction().symbol
-                    }
                 }
             }
 

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

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.ir.expressions.impl.IrSetValueImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrSuspendableExpressionImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrSuspensionPointImpl
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
+import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.symbols.IrVariableSymbol
 import org.jetbrains.kotlin.ir.symbols.impl.IrSimpleFunctionSymbolImpl
 import org.jetbrains.kotlin.ir.symbols.impl.IrVariableSymbolImpl
@@ -59,6 +60,20 @@ internal class NativeSuspendFunctionsLowering(ctx: Context): AbstractSuspendFunc
         )
     }
 
+
+    override fun IrBuilderWithScope.launchSuspendFunctionWithGivenContinuation(
+            symbol: IrSimpleFunctionSymbol, dispatchReceiver: IrExpression,
+            arguments: List<IrExpression>, continuation: IrExpression
+    ) = irCall(this@NativeSuspendFunctionsLowering.context.ir.symbols.coroutineLaunchpad).apply {
+        putValueArgument(0, irCall(symbol).apply {
+            this.dispatchReceiver = dispatchReceiver
+            arguments.forEachIndexed { index, irExpression ->
+                putValueArgument(index, irExpression)
+            }
+        })
+        putValueArgument(1, continuation)
+    }
+
     override fun buildStateMachine(stateMachineFunction: IrFunction,
                                    transformingFunction: IrFunction,
                                    argumentToPropertiesMap: Map<IrValueParameter, IrField>) {

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/optimizations/DFGBuilder.kt

@@ -613,7 +613,7 @@ internal class ModuleDFGBuilder(val context: Context, val irModule: IrModuleFrag
         private fun mapReturnType(actualType: IrType, returnType: IrType) = mapWrappedType(actualType, returnType)
 
 
-        private fun getNode(expression: IrExpression): Scoped<DataFlowIR.Node> {
+        private fun getNode(expression: IrExpression, continuationOverride: DataFlowIR.Node? = null): Scoped<DataFlowIR.Node> {
             if (expression is IrGetValue) {
                 val valueDeclaration = expression.symbol.owner
                 if (valueDeclaration is IrValueParameter)
@@ -715,7 +715,12 @@ internal class ModuleDFGBuilder(val context: Context, val irModule: IrModuleFrag
                             }
 
                             is IrCall -> when (value.symbol) {
-                                getContinuationSymbol -> getContinuation().value
+                                getContinuationSymbol -> continuationOverride ?: getContinuation().value
+
+                                symbols.coroutineLaunchpad -> getNode(
+                                        value.getValueArgument(0)!!,
+                                        continuationOverride = expressionToEdge(value.getValueArgument(1)!!).node
+                                ).value
 
                                 in arrayGetSymbols -> {
                                     val actualCallee = value.actualCallee
@@ -765,7 +770,7 @@ internal class ModuleDFGBuilder(val context: Context, val irModule: IrModuleFrag
                                             .map { expressionToEdge(it.second) }
                                             .let {
                                                 if (callee.isSuspend)
-                                                    it + DataFlowIR.Edge(getContinuation().value, null)
+                                                    it + DataFlowIR.Edge(continuationOverride ?: getContinuation().value, null)
                                                 else
                                                     it
                                             }

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/optimizations/DevirtualizationAnalysis.kt

@@ -1372,21 +1372,12 @@ internal object DevirtualizationAnalysis {
                     callSite.origin,
                     actualCallee.parentAsClass.symbol
             )
-            if (actualCallee.explicitParametersCount == arguments.size) {
+            assert(actualCallee.explicitParametersCount == arguments.size) {
-                arguments.forEachIndexed { index, argument -> call.putArgument(index, argument) }
+                "Incorrect number of arguments: expected [${actualCallee.explicitParametersCount}] but was [${arguments.size}]\n" +
-                return call
-            }
-            assert(actualCallee.isSuspend && actualCallee.explicitParametersCount == arguments.size - 1) {
-                "Incorrect number of arguments: expected [${actualCallee.explicitParametersCount}] but was [${arguments.size - 1}]\n" +
                         actualCallee.dump()
             }
-            val continuation = arguments.last()
+            arguments.forEachIndexed { index, argument -> call.putArgument(index, argument) }
-            for (index in 0..arguments.size - 2)
+            return call
-                call.putArgument(index, arguments[index])
-            return irCall(context.ir.symbols.coroutineLaunchpad, actualType).apply {
-                putValueArgument(0, call)
-                putValueArgument(1, continuation)
-            }
         }
 
         fun IrBuilderWithScope.irDevirtualizedCall(callee: IrCall, actualType: IrType,

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/optimizations/FileInitializersOptimization.kt

@@ -235,6 +235,7 @@ internal object FileInitializersOptimization {
         }
 
         private val executeImplSymbol = context.ir.symbols.executeImpl
+        private val coroutineLaunchpadSymbol = context.ir.symbols.coroutineLaunchpad
         private val getContinuationSymbol = context.ir.symbols.getContinuation
 
         private var dummySet = mutableSetOf<IrFunctionAccessExpression>()
@@ -489,6 +490,13 @@ internal object FileInitializersOptimization {
                     return curData
                 }
 
+                private fun processCoroutineLaunchpad(expression: IrCall, data: BitSet): BitSet {
+                    val call = expression.getValueArgument(0)!!
+                    val continuation = expression.getValueArgument(1)!!
+                    val curData = continuation.accept(this, data)
+                    return call.accept(this, curData)
+                }
+
                 override fun visitFunctionAccess(expression: IrFunctionAccessExpression, data: BitSet) =
                         processCall(expression, expression.actualCallee, data)
 
@@ -499,6 +507,8 @@ internal object FileInitializersOptimization {
                         return processExecuteImpl(expression, data)
                     if (expression.symbol == getContinuationSymbol)
                         return data
+                    if (expression.symbol == coroutineLaunchpadSymbol)
+                        return processCoroutineLaunchpad(expression, data)
                     if (!expression.isVirtualCall)
                         return processCall(expression, expression.actualCallee, data)
                     val devirtualizedCallSite = virtualCallSites[expression] ?: return data