JVM_IR: merge suspend view and suspend stub maps

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

@@ -115,23 +115,12 @@ class JvmBackendContext(
 
     val suspendLambdaToOriginalFunctionMap = mutableMapOf<IrFunctionReference, IrFunction>()
     val suspendFunctionOriginalToView = mutableMapOf<IrFunction, IrFunction>()
-    val suspendFunctionOriginalToStub = mutableMapOf<IrFunction, IrFunction>()
     val fakeContinuation: IrExpression = createFakeContinuation(this)
 
     val staticDefaultStubs = mutableMapOf<IrFunctionSymbol, IrFunction>()
 
     val inlineClassReplacements = MemoizedInlineClassReplacements()
 
-    internal fun recordSuspendFunctionView(function: IrFunction, view: IrFunction) {
-        val attribute = function.suspendFunctionOriginal()
-        suspendFunctionOriginalToStub.remove(attribute)
-        suspendFunctionOriginalToView[attribute] = view
-    }
-
-    internal fun recordSuspendFunctionViewStub(function: IrFunction, stub: IrFunction) {
-        suspendFunctionOriginalToStub[function.suspendFunctionOriginal()] = stub
-    }
-
     internal fun referenceClass(descriptor: ClassDescriptor): IrClassSymbol =
         symbolTable.lazyWrapper.referenceClass(descriptor)
 

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

@@ -10,7 +10,7 @@ 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.getOrCreateSuspendFunctionViewIfNeeded
+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
@@ -107,7 +107,7 @@ class IrSourceCompilerForInline(
                 function.parentAsClass.functions.find {
                     it.name.asString() == function.name.asString() + FOR_INLINE_SUFFIX &&
                             it.attributeOwnerId == (function as? IrAttributeContainer)?.attributeOwnerId
-                } ?: function.getOrCreateSuspendFunctionViewIfNeeded(classCodegen.context)
+                } ?: function.suspendFunctionView(codegen.classCodegen.context)
             else function
         val functionCodegen = object : FunctionCodegen(forInlineFunction, classCodegen, codegen.takeIf { isLambda }) {
             override fun createMethod(flags: Int, signature: JvmMethodGenericSignature): MethodVisitor {

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

@@ -560,7 +560,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                 val flag = MutableFlag(false)
                 function.accept(this, flag)
 
-                val view = function.getOrCreateSuspendFunctionViewIfNeeded(context) as IrSimpleFunction
+                val view = function.suspendFunctionView(context) as IrSimpleFunction
                 val result = mutableListOf(view)
                 if (function.body == null || !function.hasContinuation()) return result
 
@@ -621,14 +621,19 @@ 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.getOrCreateSuspendFunctionViewIfNeeded(context: JvmBackendContext): IrFunction {
+internal fun IrFunction.suspendFunctionView(context: JvmBackendContext): IrFunction {
     if (!isSuspend) return this
-    return context.suspendFunctionOriginalToView[suspendFunctionOriginal()] ?: suspendFunctionView(context)
+    val stub = suspendFunctionViewOrStub(context)
+    if (stub.body == null) {
+        val continuationParameter = stub.continuationParameter()
+        stub.body = moveBodyTo(stub, explicitParameters.zip(stub.explicitParameters.filter { it != continuationParameter }).toMap())
+    }
+    return stub
 }
 
-private fun IrFunction.getOrCreateSuspendFunctionStub(context: JvmBackendContext): IrFunction {
+internal fun IrFunction.suspendFunctionViewOrStub(context: JvmBackendContext): IrFunction {
     if (!isSuspend) return this
-    return context.suspendFunctionOriginalToStub[suspendFunctionOriginal()] ?: suspendFunctionStub(context)
+    return context.suspendFunctionOriginalToView.getOrPut(suspendFunctionOriginal()) { suspendFunctionStub(context) }
 }
 
 internal fun IrFunction.suspendFunctionOriginal(): IrFunction {
@@ -667,8 +672,6 @@ private fun IrFunction.suspendFunctionStub(context: JvmBackendContext): IrFuncti
             SUSPEND_FUNCTION_COMPLETION_PARAMETER_NAME, continuationType(context), JvmLoweredDeclarationOrigin.CONTINUATION_CLASS
         )
         function.valueParameters += valueParameters.drop(index).map { it.copyTo(function, index = it.index + 1) }
-
-        context.recordSuspendFunctionViewStub(this, function)
     }
 }
 
@@ -685,25 +688,6 @@ private fun IrFunction.continuationType(context: JvmBackendContext): IrType {
         context.ir.symbols.continuationClass.typeWith(returnType)
 }
 
-private fun IrFunction.suspendFunctionView(context: JvmBackendContext): IrFunction {
-    require(isSuspend && this is IrSimpleFunction)
-    return getOrCreateSuspendFunctionStub(context).also { function ->
-        context.recordSuspendFunctionView(this, function)
-
-        val continuationParameter =
-            function.valueParameters.find { it.origin == IrDeclarationOrigin.MASK_FOR_DEFAULT_FUNCTION }
-                ?.let { function.valueParameters[it.index - 1] } ?: function.valueParameters.last()
-
-        function.body =
-            moveBodyTo(function, explicitParameters.zip(function.explicitParameters.filter { it != continuationParameter }).toMap())
-    }
-}
-
-fun IrFunction.suspendFunctionViewOrStub(context: JvmBackendContext): IrFunction {
-    if (!isSuspend) return this
-    return context.suspendFunctionOriginalToView[suspendFunctionOriginal()] ?: getOrCreateSuspendFunctionStub(context)
-}
-
 private fun IrCall.createSuspendFunctionCallViewIfNeeded(context: JvmBackendContext, caller: IrFunction): IrCall {
     // Calls inside continuation are already generated with continuation parameter as well as calls to suspendImpls
     if (!isSuspend || caller.isInvokeSuspendOfContinuation()