Box inline class in resume path of suspend call

#KT-41429

compiler/backend/src/org/jetbrains/kotlin/codegen/ExpressionCodegen.java

@@ -46,7 +46,6 @@ import org.jetbrains.kotlin.config.ApiVersion;
 import org.jetbrains.kotlin.config.JVMAssertionsMode;
 import org.jetbrains.kotlin.config.LanguageFeature;
 import org.jetbrains.kotlin.descriptors.*;
-import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor;
 import org.jetbrains.kotlin.descriptors.impl.LocalVariableDescriptor;
 import org.jetbrains.kotlin.descriptors.impl.SyntheticFieldDescriptor;
 import org.jetbrains.kotlin.descriptors.impl.TypeAliasConstructorDescriptor;
@@ -2735,8 +2734,11 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
         callGenerator.genCall(callableMethod, resolvedCall, defaultMaskWasGenerated, this);
 
         if (maybeSuspensionPoint) {
+            // The order is important! Do not reorder!
+            // CoroutineTransformerMethodVisitor expects the marker exactly in this order.
             addReturnsUnitMarkerIfNecessary(v, resolvedCall);
             addSuspendMarker(v, false, suspensionPointKind == SuspensionPointKind.NOT_INLINE);
+            addUnboxInlineClassMarkersIfNeeded(v, resolvedCall.getResultingDescriptor(), typeMapper);
             addInlineMarker(v, false);
         }
 

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

@@ -636,6 +636,9 @@ class CoroutineCodegenForNamedFunction private constructor(
 
     override val passArityToSuperClass get() = false
 
+    private val inlineClassToBoxInInvokeSuspend: KotlinType? =
+        originalSuspendFunctionDescriptor.originalReturnTypeOfSuspendFunctionReturningUnboxedInlineClass(state.typeMapper)
+
     override fun generateBridges() {
         // Do not generate any closure bridges
     }
@@ -719,6 +722,22 @@ class CoroutineCodegenForNamedFunction private constructor(
                         callableMethod.genInvokeInstruction(codegen.v)
                     }
 
+                    if (inlineClassToBoxInInvokeSuspend != null) {
+                        with(codegen.v) {
+                            // We need to box the returned inline class in resume path.
+                            // But first, check for COROUTINE_SUSPENDED, since the function can return it
+                            dup()
+                            loadCoroutineSuspendedMarker(languageVersionSettings)
+                            val elseLabel = Label()
+                            ifacmpne(elseLabel)
+                            areturn(AsmTypes.OBJECT_TYPE)
+                            mark(elseLabel)
+                            // Now we box the inline class
+                            StackValue.coerce(AsmTypes.OBJECT_TYPE, typeMapper.mapType(inlineClassToBoxInInvokeSuspend), this)
+                            StackValue.boxInlineClass(inlineClassToBoxInInvokeSuspend, this)
+                        }
+                    }
+
                     codegen.v.visitInsn(Opcodes.ARETURN)
                 }
             }

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

@@ -114,6 +114,7 @@ class CoroutineTransformerMethodVisitor(
             if (examiner.allSuspensionPointsAreTailCalls(suspensionPoints)) {
                 examiner.replacePopsBeforeSafeUnitInstancesWithCoroutineSuspendedChecks()
                 dropSuspensionMarkers(methodNode)
+                dropUnboxInlineClassMarkers(methodNode, suspensionPoints)
                 return
             }
 
@@ -190,6 +191,7 @@ class CoroutineTransformerMethodVisitor(
         initializeFakeInlinerVariables(methodNode, stateLabels)
 
         dropSuspensionMarkers(methodNode)
+        dropUnboxInlineClassMarkers(methodNode, suspensionPoints)
         methodNode.removeEmptyCatchBlocks()
 
         if (languageVersionSettings.isReleaseCoroutines()) {
@@ -599,6 +601,17 @@ class CoroutineTransformerMethodVisitor(
         }
     }
 
+    private fun dropUnboxInlineClassMarkers(methodNode: MethodNode, suspensionPoints: List<SuspensionPoint>) {
+        for (marker in methodNode.instructions.asSequence()
+            .filter { isBeforeUnboxInlineClassMarker(it) || isAfterUnboxInlineClassMarker(it) }.toList()
+        ) {
+            methodNode.instructions.removeAll(listOf(marker.previous, marker))
+        }
+        for (suspension in suspensionPoints) {
+            methodNode.instructions.removeAll(suspension.unboxInlineClassInstructions)
+        }
+    }
+
     private fun spillVariables(suspensionPoints: List<SuspensionPoint>, methodNode: MethodNode): List<List<SpilledVariableAndField>> {
         val instructions = methodNode.instructions
         val frames =
@@ -918,6 +931,11 @@ class CoroutineTransformerMethodVisitor(
 
                 // Load continuation argument just like suspending function returns it
                 load(dataIndex, AsmTypes.OBJECT_TYPE)
+                // Unbox inline class, since this is the resume path and unlike the direct path
+                // the class is boxed.
+                for (insn in suspension.unboxInlineClassInstructions) {
+                    insn.accept(this)
+                }
 
                 visitLabel(continuationLabelAfterLoadedResult.label)
 
@@ -1117,6 +1135,16 @@ internal class SuspensionPoint(
 ) {
     lateinit var tryCatchBlocksContinuationLabel: LabelNode
 
+    val unboxInlineClassInstructions: List<AbstractInsnNode> = findUnboxInlineClassInstructions()
+
+    private fun findUnboxInlineClassInstructions(): List<AbstractInsnNode> {
+        val beforeMarker = suspensionCallEnd.next?.next ?: return emptyList()
+        if (!isBeforeUnboxInlineClassMarker(beforeMarker)) return emptyList()
+        val afterMarker = beforeMarker.findNextOrNull { isAfterUnboxInlineClassMarker(it) }
+            ?: error("Before unbox inline class marker without after unbox inline class marker")
+        return InsnSequence(beforeMarker.next, afterMarker.previous).toList()
+    }
+
     operator fun contains(insn: AbstractInsnNode): Boolean {
         for (i in InsnSequence(suspensionCallBegin, suspensionCallEnd.next)) {
             if (i == insn) return true

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

@@ -33,12 +33,14 @@ import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
 import org.jetbrains.kotlin.resolve.calls.tasks.TracingStrategy
 import org.jetbrains.kotlin.resolve.calls.tower.NewResolvedCallImpl
 import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
-import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.resolve.descriptorUtil.module
 import org.jetbrains.kotlin.resolve.descriptorUtil.resolveTopLevelClass
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes.OBJECT_TYPE
-import org.jetbrains.kotlin.types.*
+import org.jetbrains.kotlin.types.ErrorUtils
+import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.KotlinTypeFactory
+import org.jetbrains.kotlin.types.TypeConstructorSubstitution
 import org.jetbrains.kotlin.types.typeUtil.asTypeProjection
 import org.jetbrains.kotlin.types.typeUtil.makeNotNullable
 import org.jetbrains.kotlin.util.OperatorNameConventions
@@ -111,8 +113,8 @@ fun computeLabelOwner(languageVersionSettings: LanguageVersionSettings, thisName
     else
         languageVersionSettings.coroutinesJvmInternalPackageFqName().child(Name.identifier("CoroutineImpl")).topLevelClassAsmType()
 
-private val NORMALIZE_CONTINUATION_METHOD_NAME = "normalizeContinuation"
+private const val NORMALIZE_CONTINUATION_METHOD_NAME = "normalizeContinuation"
-private val GET_CONTEXT_METHOD_NAME = "getContext"
+private const val GET_CONTEXT_METHOD_NAME = "getContext"
 
 data class ResolvedCallWithRealDescriptor(val resolvedCall: ResolvedCall<*>, val fakeContinuationExpression: KtExpression)
 
@@ -140,11 +142,11 @@ fun ResolvedCall<*>.replaceSuspensionFunctionWithRealDescriptor(
     if (this is VariableAsFunctionResolvedCall) {
         val replacedFunctionCall =
             functionCall.replaceSuspensionFunctionWithRealDescriptor(project, bindingContext, isReleaseCoroutines)
-                    ?: return null
+                ?: return null
 
         @Suppress("UNCHECKED_CAST")
         return replacedFunctionCall.copy(
-            VariableAsFunctionResolvedCallImpl(
+            resolvedCall = VariableAsFunctionResolvedCallImpl(
                 replacedFunctionCall.resolvedCall as MutableResolvedCall<FunctionDescriptor>,
                 variableCall.asMutableResolvedCall(bindingContext)
             )
@@ -413,7 +415,8 @@ fun createMethodNodeForSuspendCoroutineUninterceptedOrReturn(languageVersionSett
             load(1, OBJECT_TYPE) // continuation
             checkcast(languageVersionSettings.continuationAsmType())
             invokestatic(
-                languageVersionSettings.coroutinesJvmInternalPackageFqName().child(Name.identifier("DebugProbesKt")).topLevelClassAsmType().internalName,
+                languageVersionSettings.coroutinesJvmInternalPackageFqName().child(Name.identifier("DebugProbesKt"))
+                    .topLevelClassAsmType().internalName,
                 "probeCoroutineSuspended",
                 "(${languageVersionSettings.continuationAsmType()})V",
                 false
@@ -466,7 +469,7 @@ fun FunctionDescriptor.originalReturnTypeOfSuspendFunctionReturningUnboxedInline
     val originalReturnType = originalDescriptor.returnType ?: return null
     if (!originalReturnType.isInlineClassType()) return null
     // Force boxing for primitives
-    if (AsmUtil.isPrimitive(typeMapper.mapType(originalReturnType))) return null
+    if (AsmUtil.isPrimitive(typeMapper.mapType(originalReturnType.makeNotNullable()))) return null
     // Force boxing for nullable inline class types with nullable underlying type
     if (originalReturnType.isMarkedNullable && originalReturnType.isNullableUnderlyingType()) return null
     // Force boxing if the function overrides function with return type Any
@@ -474,8 +477,6 @@ fun FunctionDescriptor.originalReturnTypeOfSuspendFunctionReturningUnboxedInline
             (it.original.returnType?.isMarkedNullable == true && it.original.returnType?.isNullableUnderlyingType() == true) ||
                     it.original.returnType?.makeNotNullable() != originalReturnType.makeNotNullable()
         }) return null
-    // TODO: Do not box `Result`
-    if (originalReturnType.constructor.declarationDescriptor?.fqNameSafe == StandardNames.RESULT_FQ_NAME) return null
     // Don't box other inline classes
     return originalReturnType
 }

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

@@ -7,9 +7,7 @@ package org.jetbrains.kotlin.codegen.inline
 
 import com.intellij.openapi.vfs.VirtualFile
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
-import org.jetbrains.kotlin.codegen.ASSERTIONS_DISABLED_FIELD_NAME
+import org.jetbrains.kotlin.codegen.*
-import org.jetbrains.kotlin.codegen.AsmUtil
-import org.jetbrains.kotlin.codegen.BaseExpressionCodegen
 import org.jetbrains.kotlin.codegen.SamWrapperCodegen.SAM_WRAPPER_SUFFIX
 import org.jetbrains.kotlin.codegen.`when`.WhenByEnumsMapping
 import org.jetbrains.kotlin.codegen.binding.CodegenBinding
@@ -17,16 +15,14 @@ import org.jetbrains.kotlin.codegen.context.CodegenContext
 import org.jetbrains.kotlin.codegen.context.CodegenContextUtil
 import org.jetbrains.kotlin.codegen.context.InlineLambdaContext
 import org.jetbrains.kotlin.codegen.context.MethodContext
+import org.jetbrains.kotlin.codegen.coroutines.originalReturnTypeOfSuspendFunctionReturningUnboxedInlineClass
 import org.jetbrains.kotlin.codegen.coroutines.unwrapInitialDescriptorForSuspendFunction
 import org.jetbrains.kotlin.codegen.optimization.common.InsnSequence
 import org.jetbrains.kotlin.codegen.optimization.common.asSequence
 import org.jetbrains.kotlin.codegen.optimization.common.intConstant
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper
-import org.jetbrains.kotlin.descriptors.ClassifierDescriptor
+import org.jetbrains.kotlin.descriptors.*
-import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
-import org.jetbrains.kotlin.descriptors.FunctionDescriptor
-import org.jetbrains.kotlin.descriptors.PackageFragmentDescriptor
 import org.jetbrains.kotlin.fileClasses.JvmFileClassUtil
 import org.jetbrains.kotlin.load.java.JvmAbi
 import org.jetbrains.kotlin.load.kotlin.VirtualFileFinder
@@ -80,6 +76,8 @@ private const val INLINE_MARKER_BEFORE_FAKE_CONTINUATION_CONSTRUCTOR_CALL = 4
 private const val INLINE_MARKER_AFTER_FAKE_CONTINUATION_CONSTRUCTOR_CALL = 5
 private const val INLINE_MARKER_BEFORE_INLINE_SUSPEND_ID = 6
 private const val INLINE_MARKER_AFTER_INLINE_SUSPEND_ID = 7
+private const val INLINE_MARKER_BEFORE_UNBOX_INLINE_CLASS = 8
+private const val INLINE_MARKER_AFTER_UNBOX_INLINE_CLASS = 9
 
 internal fun getMethodNode(
     classData: ByteArray,
@@ -413,6 +411,23 @@ 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)
+        addAfterUnboxInlineClassMarker(v)
+    }
+}
+
+private fun addBeforeUnboxInlineClassMarker(v: InstructionAdapter) {
+    v.emitInlineMarker(INLINE_MARKER_BEFORE_UNBOX_INLINE_CLASS)
+}
+
+private fun addAfterUnboxInlineClassMarker(v: InstructionAdapter) {
+    v.emitInlineMarker(INLINE_MARKER_AFTER_UNBOX_INLINE_CLASS)
+}
+
 internal fun addReturnsUnitMarkerIfNecessary(v: InstructionAdapter, resolvedCall: ResolvedCall<*>) {
     val wrapperDescriptor = resolvedCall.candidateDescriptor.safeAs<FunctionDescriptor>() ?: return
     val unsubstitutedDescriptor = wrapperDescriptor.unwrapInitialDescriptorForSuspendFunction()
@@ -477,6 +492,8 @@ internal fun isAfterInlineSuspendMarker(insn: AbstractInsnNode) = isSuspendMarke
 internal fun isReturnsUnitMarker(insn: AbstractInsnNode) = isSuspendMarker(insn, INLINE_MARKER_RETURNS_UNIT)
 internal fun isFakeContinuationMarker(insn: AbstractInsnNode) =
     insn.previous != null && isSuspendMarker(insn.previous, INLINE_MARKER_FAKE_CONTINUATION) && insn.opcode == Opcodes.ACONST_NULL
+internal fun isBeforeUnboxInlineClassMarker(insn: AbstractInsnNode) = isSuspendMarker(insn, INLINE_MARKER_BEFORE_UNBOX_INLINE_CLASS)
+internal fun isAfterUnboxInlineClassMarker(insn: AbstractInsnNode) = isSuspendMarker(insn, INLINE_MARKER_AFTER_UNBOX_INLINE_CLASS)
 
 internal fun isBeforeFakeContinuationConstructorCallMarker(insn: AbstractInsnNode) =
     isSuspendMarker(insn, INLINE_MARKER_BEFORE_FAKE_CONTINUATION_CONSTRUCTOR_CALL)

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

@@ -21,7 +21,8 @@ suspend fun signInFlowStepFirst(): Result<Unit> = try {
 
 fun box(): String {
     builder {
-        signInFlowStepFirst()
+        val res: Result<Unit> = signInFlowStepFirst()
+        if (res.exceptionOrNull()!!.message != "BOOYA") error("FAIL")
     }
     continuation!!.resumeWithException(Exception("BOOYA"))
     return "OK"