JVM_IR: Unbox inline classes in suspend functions

if the function can be verified to safely returning unboxed
inline class.
Box the return value on resume path.

compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/coroutineCodegenUtil.kt

@@ -491,7 +491,7 @@ fun InstructionAdapter.loadCoroutineSuspendedMarker(languageVersionSettings: Lan
     )
 }
 
-internal fun InstructionAdapter.generateCoroutineSuspendedCheck(languageVersionSettings: LanguageVersionSettings) {
+fun InstructionAdapter.generateCoroutineSuspendedCheck(languageVersionSettings: LanguageVersionSettings) {
     dup()
     loadCoroutineSuspendedMarker(languageVersionSettings)
     val elseLabel = Label()

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/inlineCodegenUtils.kt

@@ -32,6 +32,7 @@ import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.resolve.jvm.JvmClassName
+import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeProjectionImpl
 import org.jetbrains.kotlin.types.TypeSubstitutor
 import org.jetbrains.kotlin.util.OperatorNameConventions
@@ -414,17 +415,21 @@ fun addInlineMarker(v: InstructionAdapter, isStartNotEnd: Boolean) {
 internal fun addUnboxInlineClassMarkersIfNeeded(v: InstructionAdapter, descriptor: CallableDescriptor, typeMapper: KotlinTypeMapper) {
     val inlineClass = (descriptor as? FunctionDescriptor)?.originalReturnTypeOfSuspendFunctionReturningUnboxedInlineClass(typeMapper)
     if (inlineClass != null) {
-        addBeforeUnboxInlineClassMarker(v)
-        StackValue.unboxInlineClass(AsmTypes.OBJECT_TYPE, inlineClass, v)
-        // Suspend functions always returns Any?, but the unboxing disrupts type analysis of the bytecode.
-        // For example, if the underlying type is String, CHECKCAST String is removed.
-        // However, the unboxing is moved to the resume path, the direct path still has Any?, but now, without the CHECKCAST.
-        // Thus, we add CHECKCAST Object, which we remove, after we copy the unboxing to the resume path.
-        v.checkcast(AsmTypes.OBJECT_TYPE)
-        addAfterUnboxInlineClassMarker(v)
+        generateResumePathUnboxing(v, inlineClass)
     }
 }
 
+fun generateResumePathUnboxing(v: InstructionAdapter, inlineClass: KotlinType) {
+    addBeforeUnboxInlineClassMarker(v)
+    StackValue.unboxInlineClass(AsmTypes.OBJECT_TYPE, inlineClass, v)
+    // Suspend functions always returns Any?, but the unboxing disrupts type analysis of the bytecode.
+    // For example, if the underlying type is String, CHECKCAST String is removed.
+    // However, the unboxing is moved to the resume path, the direct path still has Any?, but now, without the CHECKCAST.
+    // Thus, we add CHECKCAST Object, which we remove, after we copy the unboxing to the resume path.
+    v.checkcast(AsmTypes.OBJECT_TYPE)
+    addAfterUnboxInlineClassMarker(v)
+}
+
 private fun addBeforeUnboxInlineClassMarker(v: InstructionAdapter) {
     v.emitInlineMarker(INLINE_MARKER_BEFORE_UNBOX_INLINE_CLASS)
 }

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

@@ -11,6 +11,8 @@ import org.jetbrains.kotlin.backend.common.lower.LocalDeclarationsLowering
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
 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.suspendFunctionOriginal
 import org.jetbrains.kotlin.codegen.ClassBuilder
 import org.jetbrains.kotlin.codegen.coroutines.CoroutineTransformerMethodVisitor
@@ -24,9 +26,8 @@ import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.IrErrorExpressionImpl
-import org.jetbrains.kotlin.ir.types.createType
+import org.jetbrains.kotlin.ir.types.*
 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
@@ -179,3 +180,23 @@ internal fun createFakeContinuation(context: JvmBackendContext): IrExpression =
     context.ir.symbols.continuationClass.createType(true, listOf(makeTypeProjection(context.irBuiltIns.anyNType, Variance.INVARIANT))),
     "FAKE_CONTINUATION"
 )
+
+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 = returnType.makeNotNull().unboxInlineClass()
+    // 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
+    // Force boxing if the function overrides function with different type modulo nullability ignoring type parameters
+    if ((this as? IrSimpleFunction)?.let {
+            it.overriddenSymbols.any { overridden ->
+                (overridden.owner.returnType.isNullable() && overridden.owner.returnType.makeNotNull().unboxInlineClass().isNullable()) ||
+                        overridden.owner.returnType.makeNotNull().classOrNull != returnType.makeNotNull().classOrNull
+            }
+        } != false) return null
+    // Don't box other inline classes
+    return returnType
+}

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

@@ -14,9 +14,11 @@ import org.jetbrains.kotlin.backend.jvm.intrinsics.JavaClassProperty
 import org.jetbrains.kotlin.backend.jvm.lower.MultifileFacadeFileEntry
 import org.jetbrains.kotlin.backend.jvm.lower.constantValue
 import org.jetbrains.kotlin.backend.jvm.lower.inlineclasses.unboxInlineClass
+import org.jetbrains.kotlin.backend.jvm.lower.suspendFunctionOriginal
 import org.jetbrains.kotlin.codegen.*
 import org.jetbrains.kotlin.codegen.AsmUtil.*
 import org.jetbrains.kotlin.codegen.coroutines.SuspensionPointKind
+import org.jetbrains.kotlin.codegen.coroutines.generateCoroutineSuspendedCheck
 import org.jetbrains.kotlin.codegen.inline.*
 import org.jetbrains.kotlin.codegen.inline.ReifiedTypeInliner.Companion.putNeedClassReificationMarker
 import org.jetbrains.kotlin.codegen.inline.ReifiedTypeInliner.OperationKind.AS
@@ -225,10 +227,20 @@ class ExpressionCodegen(
             if (irFunction.origin != JvmLoweredDeclarationOrigin.CLASS_STATIC_INITIALIZER) {
                 irFunction.markLineNumber(startOffset = irFunction is IrConstructor && irFunction.isPrimary)
             }
-            val returnType = signature.returnType
-            val returnIrType = if (irFunction !is IrConstructor) irFunction.returnType else context.irBuiltIns.unitType
-            result.materializeAt(returnType, returnIrType)
-            mv.areturn(returnType)
+            if (irFunction.isSuspend && irFunction.origin == IrDeclarationOrigin.BRIDGE) {
+                mv.areturn(OBJECT_TYPE)
+            } else {
+                var returnType = signature.returnType
+                var returnIrType = if (irFunction !is IrConstructor) irFunction.returnType else context.irBuiltIns.unitType
+                val unboxedInlineClass =
+                    irFunction.suspendFunctionOriginal().originalReturnTypeOfSuspendFunctionReturningUnboxedInlineClass()
+                if (unboxedInlineClass != null) {
+                    returnIrType = unboxedInlineClass
+                    returnType = unboxedInlineClass.asmType
+                }
+                result.materializeAt(returnType, returnIrType)
+                mv.areturn(returnType)
+            }
         }
         val endLabel = markNewLabel()
         writeLocalVariablesInTable(info, endLabel)
@@ -436,11 +448,32 @@ class ExpressionCodegen(
             if (irFunction.isInvokeSuspendOfContinuation()) IrCallGenerator.DefaultCallGenerator else callGenerator
         generatorForActualCall.genCall(callable, this, expression, isInsideCondition)
 
+        val unboxedInlineClassIrType =
+            callee.suspendFunctionOriginal().originalReturnTypeOfSuspendFunctionReturningUnboxedInlineClass()
+
         if (isSuspensionPoint != SuspensionPointKind.NEVER) {
             addSuspendMarker(mv, isStartNotEnd = false, isSuspensionPoint == SuspensionPointKind.NOT_INLINE)
+            if (unboxedInlineClassIrType != null) {
+                generateResumePathUnboxing(mv, unboxedInlineClassIrType.toIrBasedKotlinType())
+            }
             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.toIrBasedKotlinType(), mv)
+            }
+        }
+
         return when {
             (expression.type.isNothing() || expression.type.isUnit()) && irFunction.shouldContainSuspendMarkers() -> {
                 // NewInference allows casting `() -> T` to `() -> Unit`. A CHECKCAST here will fail.
@@ -461,6 +494,18 @@ class ExpressionCodegen(
                 wrapJavaClassesIntoKClasses(mv)
                 MaterialValue(this, AsmTypes.K_CLASS_ARRAY_TYPE, expression.type)
             }
+            unboxedInlineClassIrType != null && !irFunction.isInvokeSuspendOfContinuation() ->
+                object : PromisedValue(this, unboxedInlineClassIrType.asmType, unboxedInlineClassIrType) {
+                    override fun materializeAt(target: Type, irTarget: IrType) {
+                        mv.checkcast(unboxedInlineClassIrType.asmType)
+                        MaterialValue(this@ExpressionCodegen, unboxedInlineClassIrType.asmType, unboxedInlineClassIrType)
+                            .materializeAt(target, irTarget)
+                    }
+
+                    override fun discard() {
+                        pop(mv, OBJECT_TYPE)
+                    }
+                }
             else ->
                 MaterialValue(this, callable.asmMethod.returnType, callee.returnType)
         }
@@ -787,9 +832,16 @@ class ExpressionCodegen(
             return unitValue
         }
 
-        val returnType = if (owner == irFunction) signature.returnType else methodSignatureMapper.mapReturnType(owner)
+        var returnType = if (owner == irFunction) signature.returnType else methodSignatureMapper.mapReturnType(owner)
+        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, owner.returnType)
+        expression.value.accept(this, data).materializeAt(returnType, returnIrType)
         generateFinallyBlocksIfNeeded(returnType, afterReturnLabel, data)
         expression.markLineNumber(startOffset = true)
         if (isNonLocalReturn) {

compiler/testData/codegen/box/coroutines/inlineClasses/direct/returnResult.kt

@@ -1,6 +1,5 @@
 // WITH_RUNTIME
 // WITH_COROUTINES
-// IGNORE_BACKEND: JVM_IR
 // IGNORE_BACKEND_FIR: JVM_IR
 
 import helpers.*

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

@@ -1,6 +1,5 @@
 // WITH_RUNTIME
 // WITH_COROUTINES
-// IGNORE_BACKEND: JVM_IR
 // IGNORE_BACKEND_FIR: JVM_IR
 
 import helpers.*

compiler/testData/codegen/box/coroutines/inlineClasses/resumeWithException/returnResult.kt

@@ -1,6 +1,5 @@
 // WITH_RUNTIME
 // WITH_COROUTINES
-// IGNORE_BACKEND: JVM_IR
 // IGNORE_BACKEND_FIR: JVM_IR
 
 import helpers.*

compiler/testData/codegen/box/coroutines/tailCallOptimizations/returnInlineClass.kt

@@ -1,5 +1,4 @@
 // TARGET_BACKEND: JVM
-// IGNORE_BACKEND: JVM_IR
 // IGNORE_BACKEND_FIR: JVM_IR
 // FULL_JDK
 // WITH_RUNTIME

compiler/testData/codegen/bytecodeText/coroutines/inlineClasses/returnResult.kt

@@ -1,5 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
-
 @Suppress("RESULT_CLASS_IN_RETURN_TYPE")
 suspend fun signInFlowStepFirst(): Result<Unit> = Result.success(Unit)