JVM_IR: split AddContinuationLowering

One half handles suspend lambdas, the other - named functions.

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmLower.kt

@@ -336,6 +336,7 @@ private val jvmFilePhases = listOf(
     interfaceObjectCallsPhase,
 
     tailCallOptimizationPhase,
+    suspendLambdaPhase,
     addContinuationPhase,
 
     innerClassesPhase,

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

@@ -49,50 +49,71 @@ import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.utils.DFS
 import org.jetbrains.org.objectweb.asm.Type
 
+internal val suspendLambdaPhase = makeIrFilePhase(
+    ::SuspendLambdaLowering,
+    "SuspendLambda",
+    "Transform suspend lambda into continuation classes"
+)
+
 internal val addContinuationPhase = makeIrFilePhase(
     ::AddContinuationLowering,
     "AddContinuation",
-    "Add continuation classes and parameters to suspend functions and transform suspend lambdas into continuations",
+    "Add continuation classes and parameters to suspend functions",
-    prerequisite = setOf(localDeclarationsPhase, tailCallOptimizationPhase)
+    prerequisite = setOf(suspendLambdaPhase, localDeclarationsPhase, tailCallOptimizationPhase)
 )
 
-private class AddContinuationLowering(private val context: JvmBackendContext) : FileLoweringPass {
+private abstract class SuspendLoweringBase(protected val context: JvmBackendContext) : FileLoweringPass {
+    protected fun IrClass.createContinuationClassFor(
+        parent: IrDeclarationParent,
+        newOrigin: IrDeclarationOrigin,
+        newVisibility: DescriptorVisibility
+    ): IrClass = context.irFactory.buildClass {
+        name = Name.special("<Continuation>")
+        origin = newOrigin
+        visibility = newVisibility
+    }.also { irClass ->
+        irClass.createImplicitParameterDeclarationWithWrappedDescriptor()
+        irClass.superTypes += defaultType
+        irClass.parent = parent
+    }
+
+    protected fun IrClass.addFunctionOverride(function: IrSimpleFunction): IrSimpleFunction =
+        addFunction(function.name.asString(), function.returnType).apply {
+            overriddenSymbols += function.symbol
+            valueParameters += function.valueParameters.map { it.copyTo(this) }
+        }
+
+    protected fun IrClass.addFunctionOverride(
+        function: IrSimpleFunction,
+        makeBody: IrBlockBodyBuilder.(IrFunction) -> Unit
+    ): IrSimpleFunction =
+        addFunctionOverride(function).apply {
+            body = context.createIrBuilder(symbol).irBlockBody { makeBody(this@apply) }
+        }
+
+    protected fun IrSimpleFunction.generateErrorForInlineBody() {
+        val message = "This is a stub representing a copy of a suspend method without the state machine " +
+                "(used by the inliner). Since the difference is at the bytecode level, the body is " +
+                "still on the original function. Use suspendForInlineToOriginal() to retrieve it."
+        body = IrExpressionBodyImpl(startOffset, endOffset, IrErrorExpressionImpl(startOffset, endOffset, returnType, message))
+    }
+
+    protected fun IrFunction.addCompletionValueParameter(): IrValueParameter =
+        addValueParameter(SUSPEND_FUNCTION_COMPLETION_PARAMETER_NAME, continuationType())
+
+    protected fun IrFunction.continuationType(): IrType =
+        context.ir.symbols.continuationClass.typeWith(returnType).makeNullable()
+}
+
+private class SuspendLambdaLowering(context: JvmBackendContext) : SuspendLoweringBase(context) {
+    private class SuspendLambdaInfo(val reference: IrFunctionReference, val isRestricted: Boolean) {
+        val function = reference.symbol.owner
+        val arity = (reference.type as IrSimpleType).arguments.size - 1
+        val capturesCrossinline = function.valueParameters.any { reference.getValueArgument(it.index).isReadOfCrossinline() }
+        lateinit var constructor: IrConstructor
+    }
+
     override fun lower(irFile: IrFile) {
-        transformSuspendLambdasIntoContinuations(irFile)
-        // This should be done after converting lambdas into classes to avoid breaking the invariant that
-        // each lambda is referenced at most once while creating `$$forInline` methods.
-        addContinuationObjectAndContinuationParameterToSuspendFunctions(irFile)
-        addContinuationParameterToSuspendCalls(irFile)
-    }
-
-    private fun addContinuationParameterToSuspendCalls(irFile: IrFile) {
-        irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
-            val functionStack = mutableListOf<IrFunction>()
-
-            override fun visitFunction(declaration: IrFunction): IrStatement {
-                functionStack.push(declaration)
-                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 {
-                val transformed = super.visitCall(expression) as IrCall
-                return transformed.retargetToSuspendView(context, functionStack.peek() ?: return transformed) {
-                    IrCallImpl(startOffset, endOffset, type, it, origin, superQualifierSymbol)
-                }
-            }
-        })
-    }
-
-    private fun transformSuspendLambdasIntoContinuations(irFile: IrFile) {
         val inlineReferences = mutableSetOf<IrCallableReference<*>>()
         val suspendLambdas = mutableMapOf<IrFunctionReference, SuspendLambdaInfo>()
         irFile.acceptChildren(object : IrInlineReferenceLocator(context) {
@@ -117,7 +138,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
         }, null)
 
         for (lambda in suspendLambdas.values) {
-            (lambda.function.parent as IrDeclarationContainer).declarations.remove(lambda.function)
+            (lambda.function.parent as? IrDeclarationContainer)?.declarations?.remove(lambda.function)
         }
         irFile.transformChildrenVoid(object : IrElementTransformerVoidWithContext() {
             override fun visitFunctionReference(expression: IrFunctionReference): IrExpression {
@@ -261,27 +282,6 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
         }
     }
 
-    private fun IrSimpleFunction.generateErrorForInlineBody() {
-        val message = "This is a stub representing a copy of a suspend method without the state machine " +
-                "(used by the inliner). Since the difference is at the bytecode level, the body is " +
-                "still on the original function. Use suspendForInlineToOriginal() to retrieve it."
-        body = IrExpressionBodyImpl(startOffset, endOffset, IrErrorExpressionImpl(startOffset, endOffset, returnType, message))
-    }
-
-    private fun IrClass.addFunctionOverride(function: IrSimpleFunction): IrSimpleFunction =
-        addFunction(function.name.asString(), function.returnType).apply {
-            overriddenSymbols += function.symbol
-            valueParameters += function.valueParameters.map { it.copyTo(this) }
-        }
-
-    private fun IrClass.addFunctionOverride(
-        function: IrSimpleFunction,
-        makeBody: IrBlockBodyBuilder.(IrFunction) -> Unit
-    ): IrSimpleFunction =
-        addFunctionOverride(function).apply {
-            body = context.createIrBuilder(symbol).irBlockBody { makeBody(this@apply) }
-        }
-
     // Invoke function in lambdas is responsible for
     //   1) calling `create`
     //   2) starting newly created coroutine by calling `invokeSuspend`.
@@ -350,25 +350,10 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
     private fun IrBlockBodyBuilder.callInvokeSuspend(invokeSuspend: IrSimpleFunction, lambda: IrExpression): IrExpression {
         // SingletonReferencesLowering has finished a while ago, so `irUnit()` won't work anymore.
         val unitClass = context.irBuiltIns.unitClass
-        val unitField = this@AddContinuationLowering.context.cachedDeclarations.getFieldForObjectInstance(unitClass.owner)
+        val unitField = this@SuspendLambdaLowering.context.cachedDeclarations.getFieldForObjectInstance(unitClass.owner)
         return irCallOp(invokeSuspend.symbol, invokeSuspend.returnType, lambda, irGetField(null, unitField))
     }
 
-    private fun IrClass.createContinuationClassFor(
-        parent: IrDeclarationParent,
-        newOrigin: IrDeclarationOrigin,
-        newVisibility: DescriptorVisibility
-    ): IrClass = context.irFactory.buildClass {
-        name = Name.special("<Continuation>")
-        origin = newOrigin
-        visibility = newVisibility
-    }.also { irClass ->
-        irClass.createImplicitParameterDeclarationWithWrappedDescriptor()
-        irClass.superTypes += defaultType
-
-        irClass.parent = parent
-    }
-
     // Primary constructor accepts parameters equal to function reference arguments + continuation and sets the fields.
     private fun IrClass.addPrimaryConstructorForLambda(
         info: SuspendLambdaInfo,
@@ -402,12 +387,42 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                 }
             }
         }
+}
 
-    private fun IrFunction.addCompletionValueParameter(): IrValueParameter =
+private class AddContinuationLowering(context: JvmBackendContext) : SuspendLoweringBase(context) {
-        addValueParameter(SUSPEND_FUNCTION_COMPLETION_PARAMETER_NAME, continuationType())
+    override fun lower(irFile: IrFile) {
+        // This should be done after converting lambdas into classes to avoid breaking the invariant that
+        // each lambda is referenced at most once while creating `$$forInline` methods.
+        addContinuationObjectAndContinuationParameterToSuspendFunctions(irFile)
+        addContinuationParameterToSuspendCalls(irFile)
+    }
 
-    private fun IrFunction.continuationType(): IrType =
+    private fun addContinuationParameterToSuspendCalls(irFile: IrFile) {
-        context.ir.symbols.continuationClass.typeWith(returnType).makeNullable()
+        irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
+            val functionStack = mutableListOf<IrFunction>()
+
+            override fun visitFunction(declaration: IrFunction): IrStatement {
+                functionStack.push(declaration)
+                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 {
+                val transformed = super.visitCall(expression) as IrCall
+                return transformed.retargetToSuspendView(context, functionStack.peek() ?: return transformed) {
+                    IrCallImpl(startOffset, endOffset, type, it, origin, superQualifierSymbol)
+                }
+            }
+        })
+    }
 
     private fun generateContinuationClassForNamedFunction(
         irFunction: IrFunction,
@@ -650,13 +665,6 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             }
         }, null)
     }
-
-    private class SuspendLambdaInfo(val reference: IrFunctionReference, val isRestricted: Boolean) {
-        val function = reference.symbol.owner
-        val arity = (reference.type as IrSimpleType).arguments.size - 1
-        val capturesCrossinline = function.valueParameters.any { reference.getValueArgument(it.index).isReadOfCrossinline() }
-        lateinit var constructor: IrConstructor
-    }
 }
 
 // Transform `suspend fun foo(params): RetType` into `fun foo(params, $completion: Continuation<RetType>): Any?`