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.*
-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.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(
+    addFakeContinuationMarker(v)
-        Opcodes.INVOKEINTERFACE,
+
+    v.invokeinterface(
         CONTINUATION_ASM_TYPE.internalName,
         GET_CONTEXT_METHOD_NAME,
-        Type.getMethodDescriptor(COROUTINE_CONTEXT_ASM_TYPE),
+        Type.getMethodDescriptor(COROUTINE_CONTEXT_ASM_TYPE)
-        true
     )
-    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

compiler/testData/codegen/box/coroutines/coroutineContextInInlinedLambda.kt

@@ -0,0 +1,33 @@
+// WITH_RUNTIME
+// WITH_COROUTINES
+import helpers.*
+import kotlin.coroutines.experimental.*
+
+inline fun inlinedLambda(block: () -> Unit) {
+    return block()
+}
+
+suspend fun useInlined(): Boolean {
+    inlinedLambda { return coroutineContext === EmptyCoroutineContext }
+    return false
+}
+
+fun builder(c: suspend () -> Unit) {
+    c.startCoroutine(EmptyContinuation)
+}
+
+fun box(): String {
+    var res = "FAIL 1"
+    builder {
+        if (useInlined())
+            res = "OK"
+    }
+    if (res != "OK") return res
+    res = "FAIL 2"
+    builder {
+        inlinedLambda {
+            res = if (coroutineContext === EmptyCoroutineContext) "OK" else "FAIL 3"
+        }
+    }
+    return res
+}

compiler/tests-ir-jvm/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -5679,6 +5679,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             doTest(fileName);
         }
 
+        @TestMetadata("coroutineContextInInlinedLambda.kt")
+        public void testCoroutineContextInInlinedLambda() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/coroutineContextInInlinedLambda.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("coroutineToString.kt")
         public void testCoroutineToString() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/coroutineToString.kt");

compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java

@@ -5679,6 +5679,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
 
+        @TestMetadata("coroutineContextInInlinedLambda.kt")
+        public void testCoroutineContextInInlinedLambda() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/coroutineContextInInlinedLambda.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("coroutineToString.kt")
         public void testCoroutineToString() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/coroutineToString.kt");

compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -5679,6 +5679,12 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             doTest(fileName);
         }
 
+        @TestMetadata("coroutineContextInInlinedLambda.kt")
+        public void testCoroutineContextInInlinedLambda() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/coroutineContextInInlinedLambda.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("coroutineToString.kt")
         public void testCoroutineToString() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/coroutineToString.kt");

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -6285,6 +6285,12 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             doTest(fileName);
         }
 
+        @TestMetadata("coroutineContextInInlinedLambda.kt")
+        public void testCoroutineContextInInlinedLambda() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/coroutineContextInInlinedLambda.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("createCoroutineSafe.kt")
         public void testCreateCoroutineSafe() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/createCoroutineSafe.kt");

js/js.translator/src/org/jetbrains/kotlin/js/translate/intrinsic/functions/factories/CoroutineContextFIF.kt

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.js.translate.callTranslator.CallInfo
 import org.jetbrains.kotlin.js.translate.context.TranslationContext
 import org.jetbrains.kotlin.js.translate.intrinsic.functions.basic.FunctionIntrinsic
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
+import org.jetbrains.kotlin.js.translate.utils.TranslationUtils
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.inline.InlineStrategy
@@ -39,13 +40,13 @@ object CoroutineContextFIF : FunctionIntrinsicFactory {
 
     object Intrinsic : FunctionIntrinsic() {
         override fun apply(callInfo: CallInfo, arguments: List<JsExpression>, context: TranslationContext): JsExpression {
-            val continuation = context.continuationParameterDescriptor ?: error("coroutineContext called from outside of coroutine")
+            val continuation = TranslationUtils.getEnclosingContinuationParameter(context)
             val continuationDescriptor = continuation.type.constructor.declarationDescriptor as? ClassDescriptor ?: error("Continuation is not a class")
             val contContext = DescriptorUtils.getPropertyByName(continuationDescriptor.unsubstitutedMemberScope, Name.identifier("context"))
             val res = JsNameRef(
-                    context.getNameForDescriptor(contContext),
+                context.getNameForDescriptor(contContext),
-                    if (context.declarationDescriptor?.isCoroutineLambda == true) JsAstUtils.stateMachineReceiver()
+                if (context.declarationDescriptor?.isCoroutineLambda == true) JsAstUtils.stateMachineReceiver()
-                    else context.getNameForDescriptor(continuation).makeRef()
+                else TranslationUtils.translateContinuationArgument(context)
             )
             res.inlineStrategy = InlineStrategy.NOT_INLINE
             return res

js/js.translator/src/org/jetbrains/kotlin/js/translate/utils/TranslationUtils.java

@@ -354,7 +354,7 @@ public final class TranslationUtils {
     }
 
     @NotNull
-    private static VariableDescriptor getEnclosingContinuationParameter(@NotNull TranslationContext context) {
+    public static VariableDescriptor getEnclosingContinuationParameter(@NotNull TranslationContext context) {
         VariableDescriptor result = context.getContinuationParameterDescriptor();
         if (result == null) {
             assert context.getParent() != null;