JVM_IR: do not box inline classes in suspend multifile bridges

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java

@@ -10608,6 +10608,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                     runTest("compiler/testData/codegen/box/coroutines/inlineClasses/resume/invokeOperator.kt");
                 }
 
+                @Test
+                @TestMetadata("multifileBridge.kt")
+                public void testMultifileBridge() throws Exception {
+                    runTest("compiler/testData/codegen/box/coroutines/inlineClasses/resume/multifileBridge.kt");
+                }
+
                 @Test
                 @TestMetadata("overrideSuspendFun.kt")
                 public void testOverrideSuspendFun() throws Exception {

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

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
 import org.jetbrains.kotlin.backend.jvm.ir.erasedUpperBound
 import org.jetbrains.kotlin.backend.jvm.ir.isStaticInlineClassReplacement
+import org.jetbrains.kotlin.backend.jvm.lower.inlineclasses.InlineClassAbi
 import org.jetbrains.kotlin.backend.jvm.lower.inlineclasses.unboxInlineClass
 import org.jetbrains.kotlin.backend.jvm.lower.isMultifileBridge
 import org.jetbrains.kotlin.backend.jvm.lower.suspendFunctionOriginal
@@ -141,7 +142,7 @@ private fun IrFunction.isStaticInlineClassReplacementDelegatingCall(): Boolean =
             parentAsClass.declarations.find { it is IrAttributeContainer && it.attributeOwnerId == attributeOwnerId && it !== this }
                 ?.isStaticInlineClassReplacement == true
 
-internal val BRIDGE_ORIGINS = setOf(
+private val BRIDGE_ORIGINS = setOf(
     IrDeclarationOrigin.FUNCTION_FOR_DEFAULT_PARAMETER,
     JvmLoweredDeclarationOrigin.JVM_OVERLOADS_WRAPPER,
     JvmLoweredDeclarationOrigin.SYNTHETIC_ACCESSOR,
@@ -151,12 +152,17 @@ internal val BRIDGE_ORIGINS = setOf(
     IrDeclarationOrigin.BRIDGE_SPECIAL,
 )
 
-internal fun IrFunction.shouldContainSuspendMarkers(): Boolean = origin !in BRIDGE_ORIGINS &&
+// These functions contain a single `suspend` tail call, the value of which should be returned as is
+// (i.e. if it's an unboxed inline class value, it should remain unboxed).
+internal fun IrFunction.isNonBoxingSuspendDelegation(): Boolean =
+    origin in BRIDGE_ORIGINS || isMultifileBridge() || isBridgeToSuspendImplMethod()
+
+internal fun IrFunction.shouldContainSuspendMarkers(): Boolean = !isNonBoxingSuspendDelegation() &&
+        // These functions also contain a single `suspend` tail call, but if it returns an unboxed inline class value,
+        // the return of it should be checked for a suspension and potentially boxed to satisfy an interface.
         origin != IrDeclarationOrigin.DELEGATED_MEMBER &&
         !isInvokeSuspendOfContinuation() &&
-        !isMultifileBridge() &&
         !isInvokeOfSuspendCallableReference() &&
-        !isBridgeToSuspendImplMethod() &&
         !isStaticInlineClassReplacementDelegatingCall()
 
 internal fun IrFunction.hasContinuation(): Boolean = isInvokeSuspendOfLambda() ||
@@ -186,9 +192,8 @@ internal fun createFakeContinuation(context: JvmBackendContext): IrExpression =
 internal fun IrFunction.originalReturnTypeOfSuspendFunctionReturningUnboxedInlineClass(): IrType? {
     if (!isSuspend) return null
     // Check whether we in fact return inline class
-    if (returnType.classOrNull?.owner?.isInline != true) return null
+    val unboxedReturnType = InlineClassAbi.unboxType(returnType.makeNotNull()) ?: return null
-    val unboxedReturnType = returnType.makeNotNull().unboxInlineClass()
+    // Force boxing for primitives. NOTE: this also forbids unboxing a nullable inline class into a nullable primitive.
-    // Force boxing for primitives
     if (unboxedReturnType.isPrimitiveType()) return null
     // Force boxing for nullable inline class types with nullable underlying type
     if (returnType.isNullable() && unboxedReturnType.isNullable()) return null

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

@@ -229,17 +229,9 @@ class ExpressionCodegen(
             if (irFunction.origin != JvmLoweredDeclarationOrigin.CLASS_STATIC_INITIALIZER) {
                 irFunction.markLineNumber(startOffset = irFunction is IrConstructor && irFunction.isPrimary)
             }
-            val unboxedInlineClass =
+            val (returnType, returnIrType) = irFunction.returnAsmAndIrTypes()
-                irFunction.suspendFunctionOriginal().originalReturnTypeOfSuspendFunctionReturningUnboxedInlineClass()
+            result.materializeAt(returnType, returnIrType)
-            if (unboxedInlineClass != null && irFunction.origin !in BRIDGE_ORIGINS) {
+            mv.areturn(returnType)
-                result.materializeAt(unboxedInlineClass.asmType, unboxedInlineClass)
-            } else {
-                val returnIrType = if (irFunction !is IrConstructor) irFunction.returnType else context.irBuiltIns.unitType
-                result.materializeAt(signature.returnType, returnIrType)
-            }
-            // `signature.returnType` is valid here even if the return value of a suspend function was unboxed,
-            // as it's still a reference type.
-            mv.areturn(signature.returnType)
         }
         val endLabel = markNewLabel()
         writeLocalVariablesInTable(info, endLabel)
@@ -470,21 +462,6 @@ class ExpressionCodegen(
             addInlineMarker(mv, isStartNotEnd = false)
         }
 
-        if (unboxedInlineClassIrType != null) {
-            val isFunctionReference = irFunction.origin != IrDeclarationOrigin.BRIDGE &&
-                    irFunction.parentAsClass.origin == JvmLoweredDeclarationOrigin.FUNCTION_REFERENCE_IMPL
-
-            val isDelegateCall = irFunction.origin == IrDeclarationOrigin.DELEGATED_MEMBER
-
-            if (irFunction.isInvokeSuspendOfContinuation() || isFunctionReference || isDelegateCall) {
-                mv.generateCoroutineSuspendedCheck(state.languageVersionSettings)
-                mv.checkcast(unboxedInlineClassIrType.asmType)
-            }
-            if (irFunction.isInvokeSuspendOfContinuation()) {
-                StackValue.boxInlineClass(unboxedInlineClassIrType, mv, typeMapper)
-            }
-        }
-
         return when {
             (expression.type.isNothing() || expression.type.isUnit()) && irFunction.shouldContainSuspendMarkers() -> {
                 // NewInference allows casting `() -> T` to `() -> Unit`. A CHECKCAST here will fail.
@@ -505,18 +482,20 @@ class ExpressionCodegen(
                 wrapJavaClassesIntoKClasses(mv)
                 MaterialValue(this, AsmTypes.K_CLASS_ARRAY_TYPE, expression.type)
             }
-            unboxedInlineClassIrType != null && !irFunction.isInvokeSuspendOfContinuation() && irFunction.origin !in BRIDGE_ORIGINS ->
+            unboxedInlineClassIrType != null && !irFunction.isNonBoxingSuspendDelegation() -> {
-                object : PromisedValue(this, unboxedInlineClassIrType.asmType, unboxedInlineClassIrType) {
+                if (!irFunction.shouldContainSuspendMarkers()) {
-                    override fun materializeAt(target: Type, irTarget: IrType, castForReified: Boolean) {
+                    // Since the coroutine transformer won't run, we need to do this manually.
-                        mv.checkcast(unboxedInlineClassIrType.asmType)
+                    mv.generateCoroutineSuspendedCheck(state.languageVersionSettings)
-                        MaterialValue(this@ExpressionCodegen, unboxedInlineClassIrType.asmType, unboxedInlineClassIrType)
-                            .materializeAt(target, irTarget, castForReified)
-                    }
-
-                    override fun discard() {
-                        pop(mv, OBJECT_TYPE)
-                    }
                 }
+                mv.checkcast(unboxedInlineClassIrType.asmType)
+                if (irFunction.isInvokeSuspendOfContinuation()) {
+                    // TODO: why is simply materializing the value with type `Object` not enough? This branch shouldn't be needed.
+                    StackValue.boxInlineClass(unboxedInlineClassIrType, mv, typeMapper)
+                    MaterialValue(this, callable.asmMethod.returnType, callable.returnType)
+                } else {
+                    MaterialValue(this, unboxedInlineClassIrType.asmType, unboxedInlineClassIrType)
+                }
+            }
             else ->
                 MaterialValue(this, callable.asmMethod.returnType, callable.returnType)
         }
@@ -899,20 +878,25 @@ class ExpressionCodegen(
         }
     }
 
+    private fun IrFunction.returnAsmAndIrTypes(): Pair<Type, IrType> {
+        val unboxedInlineClass = suspendFunctionOriginal().originalReturnTypeOfSuspendFunctionReturningUnboxedInlineClass()
+        // In case of non-boxing delegation, the return type of the tail call was considered to be `Object`,
+        // so that's also what we'll return here to avoid casts/unboxings/etc.
+        if (unboxedInlineClass != null && !isNonBoxingSuspendDelegation()) {
+            return unboxedInlineClass.asmType to unboxedInlineClass
+        }
+        val asmType = if (this == irFunction) signature.returnType else methodSignatureMapper.mapReturnType(this)
+        val irType = if (this is IrConstructor) context.irBuiltIns.unitType else returnType
+        return asmType to irType
+    }
+
     override fun visitReturn(expression: IrReturn, data: BlockInfo): PromisedValue {
         val returnTarget = expression.returnTargetSymbol.owner
         val owner = returnTarget as? IrFunction ?: error("Unsupported IrReturnTarget: $returnTarget")
         // TODO: should be owner != irFunction
         val isNonLocalReturn = methodSignatureMapper.mapFunctionName(owner) != methodSignatureMapper.mapFunctionName(irFunction)
 
-        var returnType = if (owner == irFunction) signature.returnType else methodSignatureMapper.mapReturnType(owner)
+        val (returnType, returnIrType) = owner.returnAsmAndIrTypes()
-        var returnIrType = owner.returnType
-        val unboxedInlineClass = owner.suspendFunctionOriginal().originalReturnTypeOfSuspendFunctionReturningUnboxedInlineClass()
-        if (unboxedInlineClass != null) {
-            returnIrType = unboxedInlineClass
-            returnType = unboxedInlineClass.asmType
-        }
-
         val afterReturnLabel = Label()
         expression.value.accept(this, data).materializeAt(returnType, returnIrType)
         // In case of non-local return from suspend lambda 'materializeAt' does not box return value, box it manually.

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

@@ -261,6 +261,7 @@ private fun IrSimpleFunction.createMultifileDelegateIfNeeded(
         }
     }
 
+    function.copyAttributes(target)
     function.copyAnnotationsFrom(target)
     function.copyParameterDeclarationsFrom(target)
     function.returnType = target.returnType.substitute(target.typeParameters, function.typeParameters.map { it.defaultType })

compiler/testData/codegen/box/coroutines/inlineClasses/resume/multifileBridge.kt

@@ -0,0 +1,29 @@
+// TARGET_PLATFORM: JVM
+// WITH_RUNTIME
+// WITH_COROUTINES
+// FILE: a.kt
+@file:JvmMultifileClass
+@file:JvmName("A")
+
+import kotlin.coroutines.*
+import kotlin.coroutines.intrinsics.*
+
+@Suppress("UNSUPPORTED_FEATURE")
+inline class I(val x: Any?)
+
+suspend fun <T> suspendHere(x: T): T = suspendCoroutineUninterceptedOrReturn {
+    it.resume(x)
+    COROUTINE_SUSPENDED
+}
+
+suspend fun f(): I = I(suspendHere("OK"))
+
+// FILE: z.kt
+import helpers.*
+import kotlin.coroutines.*
+
+fun box(): String {
+    var result = "fail"
+    suspend { result = f().x as String }.startCoroutine(EmptyContinuation)
+    return result
+}

compiler/testData/codegen/box/coroutines/inlineClasses/resume/syntheticAccessor.kt

@@ -8,16 +8,13 @@ import kotlin.coroutines.intrinsics.*
 @Suppress("UNSUPPORTED_FEATURE")
 inline class I(val x: Any?)
 
-suspend fun suspendHere(): Unit = suspendCoroutineUninterceptedOrReturn { c ->
+suspend fun <T> suspendHere(x: T): T = suspendCoroutineUninterceptedOrReturn {
-    c.resume(Unit)
+    it.resume(x)
     COROUTINE_SUSPENDED
 }
 
 class C {
-    private suspend fun f(): I {
+    private suspend fun f(): I = I(suspendHere("OK"))
-        suspendHere()
-        return I("OK")
-    }
 
     fun g() = suspend { f() }
 }

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java

@@ -10608,6 +10608,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                     runTest("compiler/testData/codegen/box/coroutines/inlineClasses/resume/invokeOperator.kt");
                 }
 
+                @Test
+                @TestMetadata("multifileBridge.kt")
+                public void testMultifileBridge() throws Exception {
+                    runTest("compiler/testData/codegen/box/coroutines/inlineClasses/resume/multifileBridge.kt");
+                }
+
                 @Test
                 @TestMetadata("overrideSuspendFun.kt")
                 public void testOverrideSuspendFun() throws Exception {

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java

@@ -10608,6 +10608,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                     runTest("compiler/testData/codegen/box/coroutines/inlineClasses/resume/invokeOperator.kt");
                 }
 
+                @Test
+                @TestMetadata("multifileBridge.kt")
+                public void testMultifileBridge() throws Exception {
+                    runTest("compiler/testData/codegen/box/coroutines/inlineClasses/resume/multifileBridge.kt");
+                }
+
                 @Test
                 @TestMetadata("overrideSuspendFun.kt")
                 public void testOverrideSuspendFun() throws Exception {

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

@@ -8419,6 +8419,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                     runTest("compiler/testData/codegen/box/coroutines/inlineClasses/resume/invokeOperator.kt");
                 }
 
+                @TestMetadata("multifileBridge.kt")
+                public void testMultifileBridge() throws Exception {
+                    runTest("compiler/testData/codegen/box/coroutines/inlineClasses/resume/multifileBridge.kt");
+                }
+
                 @TestMetadata("overrideSuspendFun.kt")
                 public void testOverrideSuspendFun() throws Exception {
                     runTest("compiler/testData/codegen/box/coroutines/inlineClasses/resume/overrideSuspendFun.kt");