From 4ef2ecf9a9b2327b5c3c1b8ae1b44c0aeaa261d0 Mon Sep 17 00:00:00 2001
From: Ilmir Usmanov <ilmir.usmanov@jetbrains.com>
Date: Thu, 19 Dec 2019 19:54:57 +0100
Subject: [PATCH] JVM_IR: Add returns unit marker if suspend call returns unit

---
 .../codegen/inline/inlineCodegenUtils.kt      |  2 +-
 .../backend/jvm/codegen/ExpressionCodegen.kt  | 20 ++++++++++++++++---
 .../inlineWithoutStateMachine_ir.txt          |  2 --
 .../tailCallOptimizations/unit/simple.kt      |  1 -
 4 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/inlineCodegenUtils.kt b/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/inlineCodegenUtils.kt
index 3942c9df8c1..a2aa442efd8 100644
--- a/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/inlineCodegenUtils.kt
+++ b/compiler/backend/src/org/jetbrains/kotlin/codegen/inline/inlineCodegenUtils.kt
@@ -423,7 +423,7 @@ internal fun addReturnsUnitMarkerIfNecessary(v: InstructionAdapter, resolvedCall
     }
 }
 
-private fun addReturnsUnitMarker(v: InstructionAdapter) {
+fun addReturnsUnitMarker(v: InstructionAdapter) {
     v.emitInlineMarker(INLINE_MARKER_RETURNS_UNIT)
 }
 
diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ExpressionCodegen.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ExpressionCodegen.kt
index daf4cb788a9..092838eb97e 100644
--- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ExpressionCodegen.kt
+++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ExpressionCodegen.kt
@@ -42,6 +42,7 @@ import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
+import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes.OBJECT_TYPE
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodParameterKind
@@ -400,7 +401,7 @@ class ExpressionCodegen(
             }
             expression.symbol.descriptor is ConstructorDescriptor ->
                 throw AssertionError("IrCall with ConstructorDescriptor: ${expression.javaClass.simpleName}")
-            callee.isSuspend && !irFunction.shouldNotContainSuspendMarkers() ->
+            callee.shouldGenerateSuspendMarkers() ->
                 addInlineMarker(mv, isStartNotEnd = true)
         }
 
@@ -426,13 +427,20 @@ class ExpressionCodegen(
         expression.markLineNumber(true)
 
         // Do not generate redundant markers in continuation class.
-        if (callee.isSuspend && !irFunction.shouldNotContainSuspendMarkers()) {
+        if (callee.shouldGenerateSuspendMarkers()) {
             addSuspendMarker(mv, isStartNotEnd = true)
         }
 
         callGenerator.genCall(callable, this, expression)
 
-        if (callee.isSuspend && !irFunction.shouldNotContainSuspendMarkers()) {
+        if (callee.shouldGenerateSuspendMarkers()) {
+            // Check return type of non-lowered suspend call, in order to replace the result of the call with Unit,
+            // otherwise, it would seem like the call returns non-unit upon resume.
+            // See box/coroutines/tailCallOptimization/unit tests.
+            if (context.suspendFunctionViewToOriginal[expression.symbol.owner]?.returnType?.isUnit() == true) {
+                addReturnsUnitMarker(mv)
+            }
+
             addSuspendMarker(mv, isStartNotEnd = false)
             addInlineMarker(mv, isStartNotEnd = false)
         }
@@ -457,6 +465,12 @@ class ExpressionCodegen(
         }
     }
 
+    private fun IrFunction.shouldGenerateSuspendMarkers(): Boolean {
+        if (!isSuspend) return false
+        if (irFunction.shouldNotContainSuspendMarkers()) return false
+        return !symbol.owner.isInline || fqNameForIrSerialization == FqName("kotlin.coroutines.intrinsics.IntrinsicsKt.suspendCoroutineUninterceptedOrReturn")
+    }
+
     override fun visitVariable(declaration: IrVariable, data: BlockInfo): PromisedValue {
         val varType = typeMapper.mapType(declaration)
         val index = frameMap.enter(declaration.symbol, varType)
diff --git a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithoutStateMachine_ir.txt b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithoutStateMachine_ir.txt
index 539dacea245..d98e53f6086 100644
--- a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithoutStateMachine_ir.txt
+++ b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/inlineWithoutStateMachine_ir.txt
@@ -24,9 +24,7 @@ final class InlineWithoutStateMachineKt$complexSuspend$1 {
 }
 
 @kotlin.Metadata
-@kotlin.coroutines.jvm.internal.DebugMetadata
 final class InlineWithoutStateMachineKt$suspendHere$1 {
-    field L$0: java.lang.Object
     field label: int
     @org.jetbrains.annotations.NotNull field result: java.lang.Object
     public method <init>(p0: kotlin.coroutines.Continuation): void
diff --git a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/unit/simple.kt b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/unit/simple.kt
index edb71e5f562..609c8f90db1 100644
--- a/compiler/testData/codegen/box/coroutines/tailCallOptimizations/unit/simple.kt
+++ b/compiler/testData/codegen/box/coroutines/tailCallOptimizations/unit/simple.kt
@@ -1,6 +1,5 @@
 // IGNORE_BACKEND_FIR: JVM_IR
 // TARGET_BACKEND: JVM
-// IGNORE_BACKEND: JVM_IR
 // FULL_JDK
 // WITH_RUNTIME
 // WITH_COROUTINES