JVM_IR: keep suspend fun return type in IrCalls

Otherwise:

  * should the dispatch receiver of a call be another call to a `suspend
    fun` wrapped in something that is optimized away later, the owner of
    the method will be incorrect;

  * references to functions returning non-Unit but casted to `() ->
    Unit` (allowed by new inference) might in fact not return Unit after
    tail call optimization.

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/CoroutineCodegen.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.backend.jvm.lower.suspendFunctionOriginal
 import org.jetbrains.kotlin.codegen.ClassBuilder
 import org.jetbrains.kotlin.codegen.coroutines.CoroutineTransformerMethodVisitor
 import org.jetbrains.kotlin.codegen.coroutines.INVOKE_SUSPEND_METHOD_NAME
+import org.jetbrains.kotlin.codegen.coroutines.SUSPEND_IMPL_NAME_SUFFIX
 import org.jetbrains.kotlin.codegen.coroutines.reportSuspensionPointInsideMonitor
 import org.jetbrains.kotlin.codegen.inline.addFakeContinuationConstructorCallMarker
 import org.jetbrains.kotlin.config.isReleaseCoroutines
@@ -28,6 +29,7 @@ import org.jetbrains.kotlin.ir.types.createType
 import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
 import org.jetbrains.kotlin.ir.types.isUnit
 import org.jetbrains.kotlin.ir.util.file
+import org.jetbrains.kotlin.ir.util.functions
 import org.jetbrains.kotlin.ir.util.isSuspend
 import org.jetbrains.kotlin.ir.util.parentAsClass
 import org.jetbrains.kotlin.psi.KtElement
@@ -118,6 +120,11 @@ internal fun IrFunction.isInvokeOfSuspendCallableReference(): Boolean = isSuspen
 private fun IrFunction.isInvokeOfSuspendMainWrapper(): Boolean = !isSuspend && name.asString() == "invoke" &&
         parentAsClass.origin == JvmLoweredDeclarationOrigin.LAMBDA_IMPL
 
+private fun IrFunction.isBridgeToSuspendImplMethod(): Boolean =
+    isSuspend && this is IrSimpleFunction && parentAsClass.functions.any {
+        it.name.asString() == name.asString() + SUSPEND_IMPL_NAME_SUFFIX && it.attributeOwnerId == attributeOwnerId
+    }
+
 internal fun IrFunction.isKnownToBeTailCall(): Boolean =
     when (origin) {
         IrDeclarationOrigin.FUNCTION_FOR_DEFAULT_PARAMETER,
@@ -130,7 +137,7 @@ internal fun IrFunction.isKnownToBeTailCall(): Boolean =
         IrDeclarationOrigin.BRIDGE,
         IrDeclarationOrigin.BRIDGE_SPECIAL,
         IrDeclarationOrigin.DELEGATED_MEMBER -> true
-        else -> isInvokeOfSuspendMainWrapper() || isInvokeOfSuspendCallableReference()
+        else -> isInvokeOfSuspendMainWrapper() || isInvokeOfSuspendCallableReference() || isBridgeToSuspendImplMethod()
     }
 
 internal fun IrFunction.shouldNotContainSuspendMarkers(): Boolean =

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ExpressionCodegen.kt

@@ -436,13 +436,6 @@ class ExpressionCodegen(
         }
 
         if (expression.isSuspensionPoint()) {
-            // 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 (expression.symbol.owner.suspendFunctionOriginal().returnType.isUnit()) {
-                addReturnsUnitMarker(mv)
-            }
-
             addSuspendMarker(mv, isStartNotEnd = false)
             addInlineMarker(mv, isStartNotEnd = false)
         }
@@ -455,8 +448,12 @@ class ExpressionCodegen(
             }
             expression is IrConstructorCall ->
                 MaterialValue(this, asmType, expression.type)
-            expression.type.isUnit() -> {
+            !irFunction.shouldNotContainSuspendMarkers() && expression.type.isUnit() -> {
                 // NewInference allows casting `() -> T` to `() -> Unit`. A CHECKCAST here will fail.
+                // Also, if the callee is a suspend function with a suspending tail call, the next `resumeWith`
+                // will continue from here, but the value passed to it might not have been `Unit`. An exception
+                // is methods that do not pass through the state machine generating MethodVisitor, since getting
+                // COROUTINE_SUSPENDED here is still possible; luckily, all those methods are bridges.
                 if (callable.asmMethod.returnType != Type.VOID_TYPE)
                     MaterialValue(this, callable.asmMethod.returnType, callee.returnType).discard()
                 // don't generate redundant UNIT/pop instructions

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/FunctionCodegen.kt

@@ -11,7 +11,6 @@ import org.jetbrains.kotlin.backend.common.lower.BOUND_VALUE_PARAMETER
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
 import org.jetbrains.kotlin.codegen.AsmUtil
-import org.jetbrains.kotlin.codegen.coroutines.SUSPEND_IMPL_NAME_SUFFIX
 import org.jetbrains.kotlin.codegen.mangleNameIfNeeded
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.codegen.visitAnnotableParameterCount
@@ -142,12 +141,7 @@ open class FunctionCodegen(
             origin != IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA &&
             // This is just a template for inliner
             origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE &&
-            origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE &&
-            // Continuations are generated for suspendImpls
-            parentAsClass.functions.none {
-                it.name.asString() == name.asString() + SUSPEND_IMPL_NAME_SUFFIX &&
-                        it.attributeOwnerId == (this as? IrAttributeContainer)?.attributeOwnerId
-            }
+            origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
 
     private fun continuationClass() =
         irFunction.body!!.statements.firstIsInstance<IrClass>()

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/AddContinuationLowering.kt

@@ -75,18 +75,6 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                 return super.visitFunction(declaration).also { functionStack.pop() }
             }
 
-            override fun visitMemberAccess(expression: IrMemberAccessExpression): IrExpression {
-                val receiverType = expression.dispatchReceiver?.type
-                val newExpression = super.visitMemberAccess(expression) as IrMemberAccessExpression
-                if (receiverType != null && receiverType != newExpression.dispatchReceiver?.type) {
-                    newExpression.dispatchReceiver = IrTypeOperatorCallImpl(
-                        expression.startOffset, expression.endOffset, receiverType,
-                        IrTypeOperator.IMPLICIT_CAST, receiverType, newExpression.dispatchReceiver!!
-                    )
-                }
-                return newExpression
-            }
-
             override fun visitCall(expression: IrCall): IrExpression {
                 // This is a property, no need to add continuation parameter, since this cannot be suspend call
                 if (functionStack.isEmpty()) return super.visitCall(expression)
@@ -790,7 +778,9 @@ private fun IrCall.createSuspendFunctionCallViewIfNeeded(context: JvmBackendCont
     val view = (symbol.owner as IrSimpleFunction).suspendFunctionViewOrStub(context)
     if (view == symbol.owner) return this
 
-    return IrCallImpl(startOffset, endOffset, view.returnType, view.symbol, superQualifierSymbol = superQualifierSymbol).also {
+    // While the new callee technically returns `<original type> | COROUTINE_SUSPENDED`, the latter case is handled
+    // by a method visitor so at an IR overview we don't need to consider it.
+    return IrCallImpl(startOffset, endOffset, type, view.symbol, superQualifierSymbol = superQualifierSymbol).also {
         it.copyTypeArgumentsFrom(this)
         it.dispatchReceiver = dispatchReceiver
         it.extensionReceiver = extensionReceiver