JVM_IR: add continuations to inline suspend lambdas in the lowering

Delaying suspend view generation until the inliner is not a good idea
because it prevents lowerings after AddContinuationLowering from
further transforming the declarations.

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrInlineCodegen.kt

@@ -8,7 +8,6 @@ package org.jetbrains.kotlin.backend.jvm.codegen
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.codegen.*
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
-import org.jetbrains.kotlin.backend.jvm.lower.suspendFunctionViewOrStub
 import org.jetbrains.kotlin.backend.jvm.ir.isInlineParameter
 import org.jetbrains.kotlin.backend.jvm.ir.isLambda
 import org.jetbrains.kotlin.codegen.inline.*
@@ -222,7 +221,7 @@ class IrExpressionLambdaImpl(
     val function: IrFunction,
     private val typeMapper: IrTypeMapper,
     methodSignatureMapper: MethodSignatureMapper,
-    private val context: JvmBackendContext,
+    context: JvmBackendContext,
     isCrossInline: Boolean,
     override val isBoundCallableReference: Boolean,
     override val isExtensionLambda: Boolean
@@ -248,7 +247,7 @@ class IrExpressionLambdaImpl(
 
     override val capturedVars: List<CapturedParamDesc> = capturedParameters.keys.toList()
 
-    private val loweredMethod = methodSignatureMapper.mapAsmMethod(function.suspendFunctionViewOrStub(context))
+    private val loweredMethod = methodSignatureMapper.mapAsmMethod(function)
 
     val capturedParamsInDesc: List<Type> = if (isBoundCallableReference) {
         loweredMethod.argumentTypes.take(1)
@@ -270,9 +269,7 @@ class IrExpressionLambdaImpl(
 
     override val hasDispatchReceiver: Boolean = false
 
-    override fun getInlineSuspendLambdaViewDescriptor(): FunctionDescriptor {
+    override fun getInlineSuspendLambdaViewDescriptor(): FunctionDescriptor = function.descriptor
-        return function.suspendFunctionViewOrStub(context).descriptor
-    }
 
     override fun isCapturedSuspend(desc: CapturedParamDesc, inliningContext: InliningContext): Boolean =
         capturedParameters[desc]?.let { it.isInlineParameter() && it.type.isSuspendFunctionTypeOrSubtype() } == true

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrSourceCompilerForInline.kt

@@ -10,7 +10,6 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import org.jetbrains.kotlin.backend.common.CodegenUtil
 import org.jetbrains.kotlin.backend.common.ir.ir2string
-import org.jetbrains.kotlin.backend.jvm.lower.suspendFunctionView
 import org.jetbrains.kotlin.codegen.BaseExpressionCodegen
 import org.jetbrains.kotlin.codegen.ClassBuilder
 import org.jetbrains.kotlin.codegen.OwnerKind
@@ -100,16 +99,7 @@ class IrSourceCompilerForInline(
 
     private fun makeInlineNode(function: IrFunction, classCodegen: ClassCodegen, isLambda: Boolean): SMAPAndMethodNode {
         var node: MethodNode? = null
-        // Do not inline the generated state-machine, which was generated to support java interop of inline suspend functions.
+        val functionCodegen = object : FunctionCodegen(function, classCodegen, codegen.takeIf { isLambda }) {
-        // Instead, find its $$forInline companion (they share the same attributeOwnerId), which is generated for the inliner to use.
-        val forInlineFunction =
-            if (function.isSuspend)
-                function.parentAsClass.functions.find {
-                    it.name.asString() == function.name.asString() + FOR_INLINE_SUFFIX &&
-                            it.attributeOwnerId == (function as? IrAttributeContainer)?.attributeOwnerId
-                } ?: function.suspendFunctionView(codegen.classCodegen.context)
-            else function
-        val functionCodegen = object : FunctionCodegen(forInlineFunction, classCodegen, codegen.takeIf { isLambda }) {
             override fun createMethod(flags: Int, signature: JvmMethodGenericSignature): MethodVisitor {
                 val asmMethod = signature.asmMethod
                 node = MethodNode(
@@ -137,7 +127,16 @@ class IrSourceCompilerForInline(
         asmMethod: Method
     ): SMAPAndMethodNode {
         assert(callableDescriptor == callee.symbol.descriptor.original) { "Expected $callableDescriptor got ${callee.descriptor.original}" }
-        return makeInlineNode(callee, FakeClassCodegen(callee, codegen.classCodegen), false)
+        // Do not inline the generated state-machine, which was generated to support java interop of inline suspend functions.
+        // Instead, find its $$forInline companion (they share the same attributeOwnerId), which is generated for the inliner to use.
+        val forInlineFunction = if (callee.isSuspend)
+            callee.parentAsClass.functions.find {
+                it.name.asString() == callee.name.asString() + FOR_INLINE_SUFFIX &&
+                        it.attributeOwnerId == (callee as IrAttributeContainer).attributeOwnerId
+            } ?: callee
+        else
+            callee
+        return makeInlineNode(forInlineFunction, FakeClassCodegen(forInlineFunction, codegen.classCodegen), false)
     }
 
     override fun hasFinallyBlocks() = data.hasFinallyBlocks()

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/AddContinuationLowering.kt

@@ -28,6 +28,7 @@ import org.jetbrains.kotlin.ir.builders.declarations.*
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrFunctionReferenceImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrGetFieldImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
@@ -72,12 +73,20 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                 return super.visitFunction(declaration).also { functionStack.pop() }
             }
 
+            override fun visitFunctionReference(expression: IrFunctionReference): IrExpression {
+                val transformed = super.visitFunctionReference(expression) as IrFunctionReference
+                // The only references not yet transformed into objects are inline lambdas; the continuation
+                // for those will be taken from the inline functions they are passed to, not the enclosing scope.
+                return transformed.retargetToSuspendView(context, null) {
+                    IrFunctionReferenceImpl(startOffset, endOffset, type, it, typeArgumentsCount, reflectionTarget, origin)
+                }
+            }
+
             override fun visitCall(expression: IrCall): IrExpression {
-                // This is a property, no need to add continuation parameter, since this cannot be suspend call
+                val transformed = super.visitCall(expression) as IrCall
-                if (functionStack.isEmpty()) return super.visitCall(expression)
+                return transformed.retargetToSuspendView(context, functionStack.peek() ?: return transformed) {
-                val caller = functionStack.peek()!!
+                    IrCallImpl(startOffset, endOffset, type, it, origin, superQualifierSymbol)
-                return (super.visitCall(expression) as IrCall)
+                }
-                    .createSuspendFunctionCallViewIfNeeded(context, caller)
             }
         })
     }
@@ -548,7 +557,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
         irFile.accept(object : IrElementTransformer<MutableFlag?> {
             override fun visitClass(declaration: IrClass, data: MutableFlag?): IrStatement {
                 declaration.transformDeclarationsFlat {
-                    if (it is IrSimpleFunction && it.isSuspend && it.origin != IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA)
+                    if (it is IrSimpleFunction && it.isSuspend)
                         return@transformDeclarationsFlat transformToView(it)
                     it.accept(this, null)
                     null
@@ -560,7 +569,11 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                 val flag = MutableFlag(false)
                 function.accept(this, flag)
 
-                val view = function.suspendFunctionView(context) as IrSimpleFunction
+                val view = function.suspendFunctionViewOrStub(context) as IrSimpleFunction
+                val continuationParameter = view.continuationParameter()
+                val parameterMap = function.explicitParameters.zip(view.explicitParameters.filter { it != continuationParameter }).toMap()
+                view.body = function.moveBodyTo(view, parameterMap)
+
                 val result = mutableListOf(view)
                 if (function.body == null || !function.hasContinuation()) return result
 
@@ -621,19 +634,9 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
 
 // Transform `suspend fun foo(params): RetType` into `fun foo(params, $completion: Continuation<RetType>): Any?`
 // the result is called 'view', just to be consistent with old backend.
-internal fun IrFunction.suspendFunctionView(context: JvmBackendContext): IrFunction {
+private fun IrFunction.suspendFunctionViewOrStub(context: JvmBackendContext): IrFunction {
     if (!isSuspend) return this
-    val stub = suspendFunctionViewOrStub(context)
+    return context.suspendFunctionOriginalToView.getOrPut(suspendFunctionOriginal()) { createSuspendFunctionStub(context) }
-    if (stub.body == null) {
-        val continuationParameter = stub.continuationParameter()
-        stub.body = moveBodyTo(stub, explicitParameters.zip(stub.explicitParameters.filter { it != continuationParameter }).toMap())
-    }
-    return stub
-}
-
-internal fun IrFunction.suspendFunctionViewOrStub(context: JvmBackendContext): IrFunction {
-    if (!isSuspend) return this
-    return context.suspendFunctionOriginalToView.getOrPut(suspendFunctionOriginal()) { suspendFunctionStub(context) }
 }
 
 internal fun IrFunction.suspendFunctionOriginal(): IrFunction {
@@ -641,12 +644,12 @@ internal fun IrFunction.suspendFunctionOriginal(): IrFunction {
     return attributeOwnerId as IrFunction
 }
 
-private fun IrFunction.suspendFunctionStub(context: JvmBackendContext): IrFunction {
+private fun IrFunction.createSuspendFunctionStub(context: JvmBackendContext): IrFunction {
     require(this.isSuspend && this is IrSimpleFunction)
     return buildFunWithDescriptorForInlining(descriptor) {
-        updateFrom(this@suspendFunctionStub)
+        updateFrom(this@createSuspendFunctionStub)
-        name = this@suspendFunctionStub.name
+        name = this@createSuspendFunctionStub.name
-        origin = this@suspendFunctionStub.origin
+        origin = this@createSuspendFunctionStub.origin
         returnType = context.irBuiltIns.anyNType
     }.also { function ->
         function.parent = parent
@@ -688,31 +691,38 @@ private fun IrFunction.continuationType(context: JvmBackendContext): IrType {
         context.ir.symbols.continuationClass.typeWith(returnType)
 }
 
-private fun IrCall.createSuspendFunctionCallViewIfNeeded(context: JvmBackendContext, caller: IrFunction): IrCall {
+private fun <T : IrFunctionAccessExpression> T.retargetToSuspendView(
+    context: JvmBackendContext,
+    caller: IrFunction?,
+    copyWithTargetSymbol: T.(IrSimpleFunctionSymbol) -> T
+): T {
     // Calls inside continuation are already generated with continuation parameter as well as calls to suspendImpls
-    if (!isSuspend || caller.isInvokeSuspendOfContinuation()
+    if (!symbol.owner.isSuspend || caller?.isInvokeSuspendOfContinuation() == true
         || symbol.owner.origin == JvmLoweredDeclarationOrigin.SUSPEND_IMPL_STATIC_FUNCTION
-        || (symbol.owner.valueParameters.lastOrNull()?.name?.asString() == SUSPEND_FUNCTION_COMPLETION_PARAMETER_NAME)
+        || symbol.owner.continuationParameter() != null
     ) return this
-    val view = (symbol.owner as IrSimpleFunction).suspendFunctionViewOrStub(context)
+    val view = symbol.owner.suspendFunctionViewOrStub(context)
     if (view == symbol.owner) return this
 
     // While the new callee technically returns `<original type> | COROUTINE_SUSPENDED`, the latter case is handled
     // by a method visitor so at an IR overview we don't need to consider it.
-    return IrCallImpl(startOffset, endOffset, type, view.symbol, origin, superQualifierSymbol).also {
+    return copyWithTargetSymbol(view.symbol as IrSimpleFunctionSymbol).also {
+        it.copyAttributes(this)
         it.copyTypeArgumentsFrom(this)
         it.dispatchReceiver = dispatchReceiver
         it.extensionReceiver = extensionReceiver
         val continuationIndex = view.continuationParameter()!!.index
         for (i in 0 until valueArgumentsCount) {
-            it.putValueArgument(i + if (i >= continuationIndex) 1 else 0, getValueArgument(i)!!)
+            it.putValueArgument(i + if (i >= continuationIndex) 1 else 0, getValueArgument(i))
+        }
+        if (caller != null) {
+            // At this point the only LOCAL_FUNCTION_FOR_LAMBDAs are inline and crossinline lambdas.
+            val continuation = if (caller.origin == IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA)
+                context.fakeContinuation
+            else
+                IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, caller.continuationParameter()?.symbol
+                    ?: throw AssertionError("${caller.render()} has no continuation; can't call ${symbol.owner.render()}"))
+            it.putValueArgument(continuationIndex, continuation)
         }
-        // At this point the only LOCAL_FUNCTION_FOR_LAMBDAs are inline and crossinline lambdas
-        val continuation = if (caller.origin == IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA)
-            context.fakeContinuation
-        else
-            IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, caller.continuationParameter()?.symbol
-                ?: throw AssertionError("${caller.render()} has no continuation; can't call ${symbol.owner.render()}"))
-        it.putValueArgument(continuationIndex, continuation)
     }
 }