Fix continuaion retrieval on generated coroutineContext intrinsic code

Use fake continuation instead on real one in JVM BE.
Pass continuation parameter to closure generator in JS BE.

 #KT-22577: Fixed

compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/CoroutineTransformerMethodVisitor.kt

@@ -22,10 +22,7 @@ import org.jetbrains.kotlin.codegen.AsmUtil
 import org.jetbrains.kotlin.codegen.ClassBuilder
 import org.jetbrains.kotlin.codegen.StackValue
 import org.jetbrains.kotlin.codegen.TransformationMethodVisitor
-import org.jetbrains.kotlin.codegen.inline.MaxStackFrameSizeAndLocalsCalculator
-import org.jetbrains.kotlin.codegen.inline.isAfterSuspendMarker
-import org.jetbrains.kotlin.codegen.inline.isBeforeSuspendMarker
-import org.jetbrains.kotlin.codegen.inline.isInlineMarker
+import org.jetbrains.kotlin.codegen.inline.*
 import org.jetbrains.kotlin.codegen.optimization.DeadCodeEliminationMethodTransformer
 import org.jetbrains.kotlin.codegen.optimization.common.*
 import org.jetbrains.kotlin.codegen.optimization.fixStack.FixStackMethodTransformer
@@ -81,6 +78,13 @@ class CoroutineTransformerMethodVisitor(
             ReturnUnitMethodTransformer.transform(containingClassInternalName, methodNode)
 
             if (allSuspensionPointsAreTailCalls(containingClassInternalName, methodNode, suspensionPoints)) {
+                continuationIndex =
+                        if (isStatic(methodNode.access))
+                            Type.getArgumentTypes(methodNode.desc).size - 1
+                        else
+                            Type.getArgumentTypes(methodNode.desc).size
+                replaceFakeContinuationsWithRealOnes(methodNode, continuationIndex)
+
                 dropSuspensionMarkers(methodNode, suspensionPoints)
                 return
             }
@@ -99,6 +103,8 @@ class CoroutineTransformerMethodVisitor(
         // Actual max stack might be increased during the previous phases
         updateMaxStack(methodNode)
 
+        replaceFakeContinuationsWithRealOnes(methodNode, continuationIndex)
+
         // Remove unreachable suspension points
         // If we don't do this, then relevant frames will not be analyzed, that is unexpected from point of view of next steps (e.g. variable spilling)
         removeUnreachableSuspensionPointsAndExitPoints(methodNode, suspensionPoints)
@@ -153,6 +159,14 @@ class CoroutineTransformerMethodVisitor(
         methodNode.removeEmptyCatchBlocks()
     }
 
+    private fun replaceFakeContinuationsWithRealOnes(methodNode: MethodNode, continuationIndex: Int) {
+        val fakeContinuations = methodNode.instructions.asSequence().filter(::isFakeContinuationMarker)
+        for (fakeContinuation in fakeContinuations) {
+            methodNode.instructions.removeAll(listOf(fakeContinuation.previous.previous, fakeContinuation.previous))
+            methodNode.instructions.set(fakeContinuation, VarInsnNode(Opcodes.ALOAD, continuationIndex))
+        }
+    }
+
     private fun createInsnForReadingLabel() =
         if (isForNamedFunction)
             MethodInsnNode(

compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/coroutineCodegenUtil.kt

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.backend.common.*
 import org.jetbrains.kotlin.builtins.isBuiltinFunctionalType
 import org.jetbrains.kotlin.codegen.StackValue
 import org.jetbrains.kotlin.codegen.binding.CodegenBinding
+import org.jetbrains.kotlin.codegen.inline.addFakeContinuationMarker
 import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper
 import org.jetbrains.kotlin.codegen.topLevelClassAsmType
 import org.jetbrains.kotlin.codegen.topLevelClassInternalName
@@ -336,20 +337,21 @@ fun createMethodNodeForCoroutineContext(functionDescriptor: FunctionDescriptor):
             Opcodes.ASM5,
             Opcodes.ACC_STATIC,
             "fake",
-            Type.getMethodDescriptor(COROUTINE_CONTEXT_ASM_TYPE, CONTINUATION_ASM_TYPE),
+            Type.getMethodDescriptor(COROUTINE_CONTEXT_ASM_TYPE),
             null, null
         )
 
-    node.visitVarInsn(Opcodes.ALOAD, 0)
+    val v = InstructionAdapter(node)
 
-    node.visitMethodInsn(
-        Opcodes.INVOKEINTERFACE,
+    addFakeContinuationMarker(v)
+
+    v.invokeinterface(
         CONTINUATION_ASM_TYPE.internalName,
         GET_CONTEXT_METHOD_NAME,
-        Type.getMethodDescriptor(COROUTINE_CONTEXT_ASM_TYPE),
-        true
+        Type.getMethodDescriptor(COROUTINE_CONTEXT_ASM_TYPE)
     )
-    node.visitInsn(Opcodes.ARETURN)
+    v.areturn(COROUTINE_CONTEXT_ASM_TYPE)
+
     node.visitMaxs(1, 1)
 
     return node

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/InlineCodegen.kt

@@ -265,9 +265,6 @@ abstract class InlineCodegen<out T: BaseExpressionCodegen>(
             addInlineMarker(codegen.v, true)
         }
 
-        if (functionDescriptor.isBuiltInCoroutineContext())
-            invocationParamBuilder.addNextValueParameter(CONTINUATION_ASM_TYPE, false, continuationValue(), 0)
-
         val parameters = invocationParamBuilder.buildParameters()
 
         val info = RootInliningContext(

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/inlineCodegenUtils.kt

@@ -81,7 +81,8 @@ private const val INLINE_MARKER_FINALLY_START = "finallyStart"
 private const val INLINE_MARKER_FINALLY_END = "finallyEnd"
 private const val INLINE_MARKER_BEFORE_SUSPEND_ID = 0
 private const val INLINE_MARKER_AFTER_SUSPEND_ID = 1
-private const val INLINE_MARKET_RETURNS_UNIT = 2
+private const val INLINE_MARKER_RETURNS_UNIT = 2
+private const val INLINE_MARKER_FAKE_CONTINUATION = 3
 private val INTRINSIC_ARRAY_CONSTRUCTOR_TYPE = AsmUtil.asmTypeByClassId(classId)
 
 internal fun getMethodNode(
@@ -411,7 +412,7 @@ internal fun addSuspendMarker(v: InstructionAdapter, isStartNotEnd: Boolean) {
 }
 
 private fun addReturnsUnitMarker(v: InstructionAdapter) {
-    v.iconst(INLINE_MARKET_RETURNS_UNIT)
+    v.iconst(INLINE_MARKER_RETURNS_UNIT)
     v.visitMethodInsn(
             Opcodes.INVOKESTATIC, INLINE_MARKER_CLASS_NAME,
             "mark",
@@ -419,9 +420,26 @@ private fun addReturnsUnitMarker(v: InstructionAdapter) {
     )
 }
 
+/* There are contexts when the continuation does not yet exist, for example, in inline lambdas, which are going to
+ * be inlined into suspendable functions.
+ * In such cases we just generate the marker which is going to be replaced with real continuation on generating state machine.
+ * See [CoroutineTransformerMethodVisitor] for more info.
+ */
+internal fun addFakeContinuationMarker(v: InstructionAdapter) {
+    v.iconst(INLINE_MARKER_FAKE_CONTINUATION)
+    v.invokestatic(
+        INLINE_MARKER_CLASS_NAME,
+        "mark",
+        "(I)V", false
+    )
+    v.aconst(null)
+}
+
 internal fun isBeforeSuspendMarker(insn: AbstractInsnNode) = isSuspendMarker(insn, INLINE_MARKER_BEFORE_SUSPEND_ID)
 internal fun isAfterSuspendMarker(insn: AbstractInsnNode) = isSuspendMarker(insn, INLINE_MARKER_AFTER_SUSPEND_ID)
-internal fun isReturnsUnitMarker(insn: AbstractInsnNode) = isSuspendMarker(insn, INLINE_MARKET_RETURNS_UNIT)
+internal fun isReturnsUnitMarker(insn: AbstractInsnNode) = isSuspendMarker(insn, INLINE_MARKER_RETURNS_UNIT)
+internal fun isFakeContinuationMarker(insn: AbstractInsnNode) =
+    insn.previous != null && isSuspendMarker(insn.previous, INLINE_MARKER_FAKE_CONTINUATION) && insn.opcode == Opcodes.ACONST_NULL
 
 private fun isSuspendMarker(insn: AbstractInsnNode, id: Int) =
         isInlineMarker(insn, "mark") && insn.previous.intConstant == id