JVM_IR: inline findSuspendAndInlineLambdas

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

@@ -23,7 +23,6 @@ import org.jetbrains.kotlin.codegen.inline.coroutines.FOR_INLINE_SUFFIX
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.descriptors.Visibility
-import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.builders.*
@@ -53,8 +52,7 @@ internal val addContinuationPhase = makeIrFilePhase(
 
 private class AddContinuationLowering(private val context: JvmBackendContext) : FileLoweringPass {
     override fun lower(irFile: IrFile) {
-        val suspendLambdas = findSuspendAndInlineLambdas(irFile)
+        transformSuspendLambdasIntoContinuations(irFile)
-        transformSuspendLambdasIntoContinuations(irFile, suspendLambdas)
         // 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)
@@ -95,15 +93,28 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
         })
     }
 
-    private fun transformSuspendLambdasIntoContinuations(irFile: IrFile, suspendLambdas: List<SuspendLambdaInfo>) {
+    private fun transformSuspendLambdasIntoContinuations(irFile: IrFile) {
-        for (lambda in suspendLambdas) {
+        val inlineReferences = mutableSetOf<IrCallableReference>()
+        val suspendLambdas = mutableMapOf<IrFunctionReference, SuspendLambdaInfo>()
+        irFile.acceptChildrenVoid(object : IrInlineReferenceLocator(context) {
+            override fun visitInlineReference(argument: IrCallableReference) {
+                inlineReferences.add(argument)
+            }
+
+            override fun visitFunctionReference(expression: IrFunctionReference) {
+                expression.acceptChildrenVoid(this)
+                if (expression.isSuspend && expression.origin == IrStatementOrigin.LAMBDA && expression !in inlineReferences) {
+                    suspendLambdas[expression] = SuspendLambdaInfo(expression)
+                }
+            }
+        })
+
+        for (lambda in suspendLambdas.values) {
             (lambda.function.parent as IrDeclarationContainer).declarations.remove(lambda.function)
         }
         irFile.transformChildrenVoid(object : IrElementTransformerVoidWithContext() {
             override fun visitFunctionReference(expression: IrFunctionReference): IrExpression {
-                if (!expression.isSuspend)
+                val info = suspendLambdas[expression] ?: return super.visitFunctionReference(expression)
-                    return expression
-                val info = suspendLambdas.singleOrNull { it.function == expression.symbol.owner } ?: return expression
                 return context.createIrBuilder(expression.symbol, expression.startOffset, expression.endOffset).run {
                     val expressionArguments = expression.getArguments().map { it.second }
                     irBlock {
@@ -117,7 +128,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                             "Inconsistency between callable reference to suspend lambda and the corresponding continuation"
                         }
                         +irCall(constructor.symbol).apply {
-                            for (typeParameter in info.constructor.parentAsClass.typeParameters) {
+                            for (typeParameter in constructor.parentAsClass.typeParameters) {
                                 putTypeArgument(typeParameter.index, expression.getTypeArgument(typeParameter.index))
                             }
                             expressionArguments.forEachIndexed { index, argument ->
@@ -664,25 +675,6 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
         })
     }
 
-    private fun findSuspendAndInlineLambdas(irElement: IrElement): List<SuspendLambdaInfo> {
-        val suspendLambdas = arrayListOf<SuspendLambdaInfo>()
-        val inlineReferences = mutableSetOf<IrCallableReference>()
-        irElement.acceptChildrenVoid(object : IrInlineReferenceLocator(context) {
-            override fun visitInlineReference(argument: IrCallableReference) {
-                inlineReferences.add(argument)
-            }
-
-            override fun visitFunctionReference(expression: IrFunctionReference) {
-                expression.acceptChildrenVoid(this)
-
-                if (expression.isSuspend && expression.origin == IrStatementOrigin.LAMBDA && expression !in inlineReferences) {
-                    suspendLambdas += SuspendLambdaInfo(expression)
-                }
-            }
-        })
-        return suspendLambdas
-    }
-
     private class SuspendLambdaInfo(val reference: IrFunctionReference) {
         val function = reference.symbol.owner
         val arity = (reference.type as IrSimpleType).arguments.size - 1