JS: fix translation of safe calls to suspend unit functions

See KT-21317

compiler/testData/codegen/box/coroutines/unitTypeReturn/unitSafeCall.kt

@@ -0,0 +1,53 @@
+// WITH_RUNTIME
+// WITH_COROUTINES
+import helpers.*
+import kotlin.coroutines.experimental.*
+import kotlin.coroutines.experimental.intrinsics.*
+
+var log = ""
+var postponed: () -> Unit = { }
+var complete = false
+
+suspend fun suspendHere(x: String): Unit {
+    log += "suspendHere($x);"
+    return suspendCoroutineOrReturn { c ->
+        postponed = { c.resume(Unit) }
+        log += "suspended;"
+        COROUTINE_SUSPENDED
+    }
+}
+
+class A(val x: String) {
+    suspend fun foo() = suspendHere("A.foo($x)")
+}
+
+fun builder(c: suspend () -> Unit) {
+    c.startCoroutine(handleResultContinuation {
+        log += "complete;"
+        complete = true
+    })
+}
+
+fun box(): String {
+    val a1: A? = A("*")
+    val a2: A? = null
+    builder {
+        val r1 = a1?.foo().simpleName + ";"
+        log += r1
+        val r2 = a2?.foo().simpleName + ";"
+        log += r2
+        A("@").foo()
+    }
+
+    while (!complete) {
+        log += "resuming;"
+        postponed()
+    }
+
+    if (log != "suspendHere(A.foo(*));suspended;resuming;Unit;null;suspendHere(A.foo(@));suspended;resuming;complete;") return "fail: $log"
+
+    return "OK"
+}
+
+val Any?.simpleName: String
+    get() = if (this == null) "null" else "Unit"

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

@@ -6396,6 +6396,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/unitTypeReturn/suspendReturn.kt");
                 doTest(fileName);
             }
+
+            @TestMetadata("unitSafeCall.kt")
+            public void testUnitSafeCall() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/unitTypeReturn/unitSafeCall.kt");
+                doTest(fileName);
+            }
         }
 
         @TestMetadata("compiler/testData/codegen/box/coroutines/varSpilling")

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

@@ -6396,6 +6396,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/unitTypeReturn/suspendReturn.kt");
                 doTest(fileName);
             }
+
+            @TestMetadata("unitSafeCall.kt")
+            public void testUnitSafeCall() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/unitTypeReturn/unitSafeCall.kt");
+                doTest(fileName);
+            }
         }
 
         @TestMetadata("compiler/testData/codegen/box/coroutines/varSpilling")

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

@@ -6396,6 +6396,12 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/unitTypeReturn/suspendReturn.kt");
                 doTest(fileName);
             }
+
+            @TestMetadata("unitSafeCall.kt")
+            public void testUnitSafeCall() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/unitTypeReturn/unitSafeCall.kt");
+                doTest(fileName);
+            }
         }
 
         @TestMetadata("compiler/testData/codegen/box/coroutines/varSpilling")

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

@@ -6954,6 +6954,12 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/unitTypeReturn/suspendReturn.kt");
                 doTest(fileName);
             }
+
+            @TestMetadata("unitSafeCall.kt")
+            public void testUnitSafeCall() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/unitTypeReturn/unitSafeCall.kt");
+                doTest(fileName);
+            }
         }
 
         @TestMetadata("compiler/testData/codegen/box/coroutines/varSpilling")

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

@@ -20,16 +20,15 @@ import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.diagnostics.DiagnosticUtils
-import org.jetbrains.kotlin.js.backend.ast.JsBlock
+import org.jetbrains.kotlin.js.backend.ast.*
-import org.jetbrains.kotlin.js.backend.ast.JsConditional
-import org.jetbrains.kotlin.js.backend.ast.JsExpression
-import org.jetbrains.kotlin.js.backend.ast.JsNullLiteral
 import org.jetbrains.kotlin.js.backend.ast.metadata.type
 import org.jetbrains.kotlin.js.translate.context.TranslationContext
 import org.jetbrains.kotlin.js.translate.reference.CallArgumentTranslator
 import org.jetbrains.kotlin.js.translate.reference.ReferenceTranslator
+import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
 import org.jetbrains.kotlin.js.translate.utils.JsDescriptorUtils.getReceiverParameterForReceiver
 import org.jetbrains.kotlin.js.translate.utils.TranslationUtils
+import org.jetbrains.kotlin.js.translate.utils.createCoroutineResult
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.calls.callUtil.isSafeCall
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
@@ -45,6 +44,8 @@ interface CallInfo {
     val extensionReceiver: JsExpression?
 
     fun constructSafeCallIfNeeded(result: JsExpression): JsExpression
+
+    fun constructSuspendSafeCallIfNeeded(result: JsStatement): JsStatement
 }
 
 abstract class AbstractCallInfo : CallInfo {
@@ -165,7 +166,7 @@ private fun TranslationContext.createCallInfo(
 
     if (dispatchReceiverType != null) {
         dispatchReceiver = dispatchReceiver?.let {
-            TranslationUtils.coerce(this, it, dispatchReceiverType!!)
+            TranslationUtils.coerce(this, it, dispatchReceiverType)
         }
     }
 
@@ -193,5 +194,29 @@ private fun TranslationContext.createCallInfo(
                 notNullConditionalForSafeCall
             }
         }
+
+        override fun constructSuspendSafeCallIfNeeded(result: JsStatement): JsStatement {
+            return if (notNullConditionalForSafeCall == null) {
+                result
+            }
+            else {
+                val callElement = resolvedCall.call.callElement
+                val coroutineResult = context.createCoroutineResult(resolvedCall)
+                val nullAssignment = JsAstUtils.assignment(coroutineResult, JsNullLiteral()).source(callElement)
+
+                val thenBlock = JsBlock()
+                thenBlock.statements += result
+                val thenContext = context.innerBlock(thenBlock)
+                val lhs = coroutineResult.deepCopy()
+                val rhsOriginal = coroutineResult.deepCopy().apply { type = resolvedCall.getReturnType() }
+                val rhs = TranslationUtils.coerce(thenContext, rhsOriginal, resolvedCall.getReturnType().makeNullable())
+                if (rhs != rhsOriginal) {
+                    thenBlock.statements += JsAstUtils.asSyntheticStatement(JsAstUtils.assignment(lhs, rhs).source(callElement))
+                }
+
+                val thenStatement = if (thenBlock.statements.size == 1) thenBlock.statements.first() else thenBlock
+                JsIf(notNullConditionalForSafeCall.testExpression, thenStatement, nullAssignment.makeStmt())
+            }
+        }
     }
 }

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

@@ -22,7 +22,6 @@ import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.VariableDescriptor
 import org.jetbrains.kotlin.js.backend.ast.JsExpression
-import org.jetbrains.kotlin.js.backend.ast.JsNameRef
 import org.jetbrains.kotlin.js.backend.ast.metadata.*
 import org.jetbrains.kotlin.js.translate.context.TranslationContext
 import org.jetbrains.kotlin.js.translate.general.Translation
@@ -59,7 +58,8 @@ object CallTranslator {
                      extensionOrDispatchReceiver: JsExpression? = null
     ): JsExpression {
         val variableAccessInfo = VariableAccessInfo(context.getCallInfo(resolvedCall, extensionOrDispatchReceiver), null)
-        val result = variableAccessInfo.translateVariableAccess().source(resolvedCall.call.callElement)
+        val result = variableAccessInfo.constructSafeCallIfNeeded(variableAccessInfo.translateVariableAccess())
+                .source(resolvedCall.call.callElement)
         result.type = TranslationUtils.getReturnTypeForCoercion(resolvedCall.resultingDescriptor.original)
         return result
     }
@@ -72,7 +72,8 @@ object CallTranslator {
         val type = TranslationUtils.getReturnTypeForCoercion(resolvedCall.resultingDescriptor)
         val coerceValue = TranslationUtils.coerce(context, value, type)
         val variableAccessInfo = VariableAccessInfo(context.getCallInfo(resolvedCall, extensionOrDispatchReceiver), coerceValue)
-        val result = variableAccessInfo.translateVariableAccess().source(resolvedCall.call.callElement)
+        val result = variableAccessInfo.constructSafeCallIfNeeded(variableAccessInfo.translateVariableAccess())
+                .source(resolvedCall.call.callElement)
         result.type = context.currentModule.builtIns.unitType
         return result
     }
@@ -108,7 +109,7 @@ private fun translateCall(
         assert(explicitReceivers.extensionReceiver == null) { "VariableAsFunctionResolvedCall must have one receiver" }
         val variableCall = resolvedCall.variableCall
 
-        val result = if (variableCall.expectedReceivers()) {
+        return if (variableCall.expectedReceivers()) {
             val newReceiver = CallTranslator.translateGet(context, variableCall, explicitReceivers.extensionOrDispatchReceiver)
             translateFunctionCall(context, resolvedCall.functionCall, resolvedCall.variableCall, ExplicitReceivers(newReceiver))
         } else {
@@ -123,8 +124,6 @@ private fun translateCall(
                                       ExplicitReceivers(dispatchReceiver, explicitReceivers.extensionOrDispatchReceiver))
             }
         }
-
-        return result
     }
 
     val call = resolvedCall.call
@@ -145,7 +144,8 @@ private fun translateFunctionCall(
     val rangeCheck = RangeCheckTranslator(context).translateAsRangeCheck(resolvedCall, explicitReceivers)
     if (rangeCheck != null) return rangeCheck
 
-    val callExpression = context.getCallInfo(resolvedCall, explicitReceivers).translateFunctionCall()
+    val callInfo = context.getCallInfo(resolvedCall, explicitReceivers)
+    var callExpression = callInfo.translateFunctionCall()
 
     if (CallExpressionTranslator.shouldBeInlined(inlineResolvedCall.resultingDescriptor, context)) {
         setInlineCallMetadata(callExpression, resolvedCall.call.callElement as KtExpression,
@@ -154,21 +154,20 @@ private fun translateFunctionCall(
 
     if (resolvedCall.resultingDescriptor.isSuspend) {
         if (context.isInStateMachine) {
-            context.currentBlock.statements += JsAstUtils.asSyntheticStatement(callExpression.apply {
+            val statement = callInfo.constructSuspendSafeCallIfNeeded(JsAstUtils.asSyntheticStatement(callExpression.apply {
                 isSuspend = true
                 source = resolvedCall.call.callElement
-            })
+            }))
-            val coroutineRef = TranslationUtils.translateContinuationArgument(context).apply { source = resolvedCall.call.callElement }
+            context.currentBlock.statements += statement
-            return context.defineTemporary(JsNameRef("\$\$coroutineResult\$\$", coroutineRef).apply {
+            return context.createCoroutineResult(resolvedCall)
-                sideEffects = SideEffectKind.DEPENDS_ON_STATE
-                source = resolvedCall.call.callElement
-                coroutineResult = true
-            })
         }
         else {
             callExpression.isTailCallSuspend = true
         }
     }
+    else {
+        callExpression = callInfo.constructSafeCallIfNeeded(callExpression)
+    }
 
     callExpression.type = resolvedCall.getReturnType().let { if (resolvedCall.call.isSafeCall()) it.makeNullable() else it }
     mayBeMarkByRangeMetadata(resolvedCall, callExpression)
@@ -254,7 +253,7 @@ abstract class CallCase<in I : CallInfo> {
     protected open fun I.bothReceivers(): JsExpression = unsupported()
 
     fun translate(callInfo: I): JsExpression {
-        val result = if (callInfo.dispatchReceiver == null) {
+        return if (callInfo.dispatchReceiver == null) {
             if (callInfo.extensionReceiver == null)
                 callInfo.noReceivers()
             else
@@ -265,8 +264,6 @@ abstract class CallCase<in I : CallInfo> {
             } else
                 callInfo.bothReceivers()
         }
-
-        return callInfo.constructSafeCallIfNeeded(result)
     }
 }
 
@@ -279,20 +276,7 @@ interface DelegateIntrinsic<in I : CallInfo> {
     fun I.getDescriptor(): CallableDescriptor
     fun I.getArgs(): List<JsExpression>
 
-    fun intrinsic(callInfo: I): JsExpression? {
+    fun intrinsic(callInfo: I): JsExpression? = if (callInfo.canBeApply()) callInfo.getIntrinsic() else null
-        val result =
-                if (callInfo.canBeApply()) {
-                    callInfo.getIntrinsic()
-                } else {
-                    null
-                }
-
-        return if (result != null) {
-            callInfo.constructSafeCallIfNeeded(result)
-        } else {
-            null
-        }
-    }
 
     private fun I.getIntrinsic(): JsExpression? {
         val descriptor = getDescriptor()

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

@@ -27,9 +27,7 @@ import org.jetbrains.kotlin.descriptors.TypeParameterDescriptor
 import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.js.backend.ast.*
-import org.jetbrains.kotlin.js.backend.ast.metadata.CoroutineMetadata
+import org.jetbrains.kotlin.js.backend.ast.metadata.*
-import org.jetbrains.kotlin.js.backend.ast.metadata.coroutineMetadata
-import org.jetbrains.kotlin.js.backend.ast.metadata.exportedPackage
 import org.jetbrains.kotlin.js.translate.context.Namer
 import org.jetbrains.kotlin.js.translate.context.TranslationContext
 import org.jetbrains.kotlin.js.translate.intrinsic.functions.basic.FunctionIntrinsicWithReceiverComputed
@@ -41,6 +39,7 @@ import org.jetbrains.kotlin.psi.KtFunctionLiteral
 import org.jetbrains.kotlin.psi.KtLambdaExpression
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DescriptorUtils
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.descriptorUtil.hasOrInheritsParametersWithDefaultValue
 import org.jetbrains.kotlin.resolve.source.getPsi
 import org.jetbrains.kotlin.types.KotlinType
@@ -243,4 +242,15 @@ fun createPrototypeStatements(superName: JsName, name: JsName): List<JsStatement
     val constructorStatement = JsAstUtils.assignment(constructorRef, classRef.deepCopy()).makeStmt()
 
     return listOf(prototypeStatement, constructorStatement)
+}
+
+fun TranslationContext.createCoroutineResult(resolvedCall: ResolvedCall<*>): JsExpression {
+    val callElement = resolvedCall.call.callElement
+    val coroutineRef = TranslationUtils.translateContinuationArgument(this).source(callElement)
+    return JsNameRef("\$\$coroutineResult\$\$", coroutineRef).apply {
+        sideEffects = SideEffectKind.DEPENDS_ON_STATE
+        source = callElement
+        coroutineResult = true
+        synthetic = true
+    }
 }