JS: fix suspend lambda inlining

* Also avoid casts when setting inline metadata

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

@@ -893,6 +893,11 @@ public class BoxJsTestGenerated extends AbstractBoxJsTest {
             runTest("js/js.translator/testData/box/coroutines/debugStatement.kt");
         }
 
+        @TestMetadata("inlineLambdaNonLocalReturn.kt")
+        public void testInlineLambdaNonLocalReturn() throws Exception {
+            runTest("js/js.translator/testData/box/coroutines/inlineLambdaNonLocalReturn.kt");
+        }
+
         @TestMetadata("localVarOptimization.kt")
         public void testLocalVarOptimization() throws Exception {
             runTest("js/js.translator/testData/box/coroutines/localVarOptimization.kt");

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

@@ -893,6 +893,11 @@ public class IrBoxJsTestGenerated extends AbstractIrBoxJsTest {
             runTest("js/js.translator/testData/box/coroutines/debugStatement.kt");
         }
 
+        @TestMetadata("inlineLambdaNonLocalReturn.kt")
+        public void testInlineLambdaNonLocalReturn() throws Exception {
+            runTest("js/js.translator/testData/box/coroutines/inlineLambdaNonLocalReturn.kt");
+        }
+
         @TestMetadata("localVarOptimization.kt")
         public void testLocalVarOptimization() throws Exception {
             runTest("js/js.translator/testData/box/coroutines/localVarOptimization.kt");

js/js.translator/src/org/jetbrains/kotlin/js/translate/callTranslator/CallTranslator.kt

@@ -18,6 +18,8 @@ package org.jetbrains.kotlin.js.translate.callTranslator
 
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.builtins.isFunctionTypeOrSubtype
+import org.jetbrains.kotlin.builtins.isSuspendFunctionType
+import org.jetbrains.kotlin.builtins.isSuspendFunctionTypeOrSubtype
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.VariableDescriptor
@@ -109,7 +111,7 @@ private fun translateCall(
     if (resolvedCall is VariableAsFunctionResolvedCall) {
         assert(explicitReceivers.extensionReceiver == null) { "VariableAsFunctionResolvedCall must have one receiver" }
         val variableCall = resolvedCall.variableCall
-        val isFunctionType = variableCall.resultingDescriptor.type.isFunctionTypeOrSubtype
+        val isFunctionType = variableCall.resultingDescriptor.type.run { isFunctionTypeOrSubtype || isSuspendFunctionTypeOrSubtype }
         val inlineCall = if (isFunctionType) variableCall else resolvedCall
 
         val newExplicitReceivers = if (variableCall.expectedReceivers()) {
@@ -149,9 +151,7 @@ private fun translateFunctionCall(
     var callExpression = callInfo.translateFunctionCall()
 
     if (CallExpressionTranslator.shouldBeInlined(inlineResolvedCall.resultingDescriptor, context)) {
-        val callElement = resolvedCall.call.callElement
-        val ktExpression = (callElement as? KtWhenConditionInRange)?.rangeExpression ?: callElement as KtExpression
-        setInlineCallMetadata(callExpression, ktExpression, inlineResolvedCall.resultingDescriptor, context)
+        setInlineCallMetadata(callExpression, resolvedCall.call.callElement, inlineResolvedCall.resultingDescriptor, context)
     }
 
     if (resolvedCall.resultingDescriptor.isSuspend) {
@@ -171,7 +171,6 @@ private fun translateFunctionCall(
     return callExpression
 }
 
-
 private fun mayBeMarkByRangeMetadata(resolvedCall: ResolvedCall<out FunctionDescriptor>, callExpression: JsExpression) {
     when (resolvedCall.resultingDescriptor.fqNameSafe) {
         intRangeToFqName -> {

js/js.translator/src/org/jetbrains/kotlin/js/translate/utils/inlineUtils.kt

@@ -24,7 +24,7 @@ import org.jetbrains.kotlin.js.backend.ast.metadata.*
 import org.jetbrains.kotlin.js.inline.util.isCallInvocation
 import org.jetbrains.kotlin.js.translate.context.TranslationContext
 import org.jetbrains.kotlin.js.translate.reference.CallExpressionTranslator
-import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.inline.InlineStrategy
 
@@ -40,10 +40,10 @@ import org.jetbrains.kotlin.resolve.inline.InlineStrategy
  *  (x != null) ? fn.call(x, y) : null
  */
 fun setInlineCallMetadata(
-        expression: JsExpression,
-        psiElement: KtExpression,
-        descriptor: CallableDescriptor,
-        context: TranslationContext
+    expression: JsExpression,
+    psiElement: KtElement,
+    descriptor: CallableDescriptor,
+    context: TranslationContext
 ) {
     assert(CallExpressionTranslator.shouldBeInlined(descriptor)) {
         "Expected descriptor of callable, that should be inlined, but got: $descriptor"
@@ -70,14 +70,14 @@ fun setInlineCallMetadata(
 
 fun setInlineCallMetadata(
         expression: JsExpression,
-        psiElement: KtExpression,
+        psiElement: KtElement,
         resolvedCall: ResolvedCall<*>,
         context: TranslationContext
 ) = setInlineCallMetadata(expression, psiElement, PsiUtils.getFunctionDescriptor(resolvedCall), context)
 
 fun setInlineCallMetadata(
         nameRef: JsNameRef,
-        psiElement: KtExpression,
+        psiElement: KtElement,
         descriptor: CallableDescriptor,
         context: TranslationContext
 ) {

js/js.translator/testData/box/coroutines/inlineLambdaNonLocalReturn.kt

@@ -0,0 +1,27 @@
+// EXPECTED_REACHABLE_NODES: 1305
+// IGNORE_BACKEND: JS_IR
+
+import kotlin.coroutines.*
+
+suspend inline fun doTwice(block: suspend () -> Unit) {
+    block()
+    block()
+}
+
+var testResult: String = ""
+
+fun build(c: suspend () -> Unit) {
+    c.startCoroutine(Continuation<Unit>(EmptyCoroutineContext) { })
+}
+
+fun box(): String {
+
+    build {
+        doTwice {
+            testResult +=  "OK"
+            return@build
+        }
+    }
+
+    return testResult
+}