JVM_IR: do not place suspend markers around crossinline lambda calls

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

@@ -8,7 +8,6 @@ package org.jetbrains.kotlin.codegen.inline
 import com.intellij.util.ArrayUtil
 import org.jetbrains.kotlin.codegen.*
 import org.jetbrains.kotlin.codegen.coroutines.DEBUG_METADATA_ANNOTATION_ASM_TYPE
-import org.jetbrains.kotlin.codegen.coroutines.isCapturedSuspendLambda
 import org.jetbrains.kotlin.codegen.coroutines.isCoroutineSuperClass
 import org.jetbrains.kotlin.codegen.coroutines.isResumeImplMethodName
 import org.jetbrains.kotlin.codegen.inline.coroutines.CoroutineTransformer
@@ -526,10 +525,8 @@ class AnonymousObjectTransformer(
                         alreadyAddedParam?.newFieldName ?: getNewFieldName(desc.fieldName, false),
                         alreadyAddedParam != null
                     )
-                    if (info is PsiExpressionLambda && info.closure.captureVariables.any { it.value.fieldName == desc.fieldName }) {
+                    if (info is ExpressionLambda && info.isCapturedSuspend(desc, inliningContext)) {
-                        recapturedParamInfo.functionalArgument = NonInlineableArgumentForInlineableParameterCalledInSuspend(
+                        recapturedParamInfo.functionalArgument = NonInlineableArgumentForInlineableParameterCalledInSuspend
-                            isCapturedSuspendLambda(info.closure, desc.fieldName, inliningContext.state.bindingContext)
-                        )
                     }
                     val composed = StackValue.field(
                         desc.type,

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

@@ -391,7 +391,7 @@ abstract class InlineCodegen<out T : BaseExpressionCodegen>(
                 info = invocationParamBuilder.addNextValueParameter(jvmType, false, remappedValue, parameterIndex)
                 info.functionalArgument = when (kind) {
                     ValueKind.NON_INLINEABLE_ARGUMENT_FOR_INLINE_PARAMETER_CALLED_IN_SUSPEND ->
-                        NonInlineableArgumentForInlineableParameterCalledInSuspend(kotlinType?.isSuspendFunctionTypeOrSubtype == true)
+                        NonInlineableArgumentForInlineableParameterCalledInSuspend
                     ValueKind.NON_INLINEABLE_ARGUMENT_FOR_INLINE_SUSPEND_PARAMETER -> NonInlineableArgumentForInlineableSuspendParameter
                     else -> null
                 }

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

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.codegen.binding.CodegenBinding.*
 import org.jetbrains.kotlin.codegen.binding.MutableClosure
 import org.jetbrains.kotlin.codegen.context.EnclosedValueDescriptor
 import org.jetbrains.kotlin.codegen.coroutines.getOrCreateJvmSuspendFunctionView
+import org.jetbrains.kotlin.codegen.coroutines.isCapturedSuspendLambda
 import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.config.isReleaseCoroutines
@@ -82,7 +83,7 @@ abstract class LambdaInfo(@JvmField val isCrossInline: Boolean) : FunctionalArgu
     }
 }
 
-class NonInlineableArgumentForInlineableParameterCalledInSuspend(val isSuspend: Boolean) : FunctionalArgument
+object NonInlineableArgumentForInlineableParameterCalledInSuspend : FunctionalArgument
 object NonInlineableArgumentForInlineableSuspendParameter : FunctionalArgument
 
 
@@ -215,6 +216,7 @@ abstract class ExpressionLambda(isCrossInline: Boolean) : LambdaInfo(isCrossInli
     }
 
     abstract fun getInlineSuspendLambdaViewDescriptor(): FunctionDescriptor
+    abstract fun isCapturedSuspend(desc: CapturedParamDesc, inliningContext: InliningContext): Boolean
 }
 
 class PsiExpressionLambda(
@@ -328,4 +330,7 @@ class PsiExpressionLambda(
             typeMapper.bindingContext
         )
     }
+
+    override fun isCapturedSuspend(desc: CapturedParamDesc, inliningContext: InliningContext): Boolean =
+        isCapturedSuspendLambda(closure, desc.fieldName, inliningContext.state.bindingContext)
 }

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

@@ -156,7 +156,8 @@ class PsiInlineCodegen(
         } else {
             val value = codegen.gen(argumentExpression)
             val kind = when {
-                isCallSiteIsSuspend(valueParameterDescriptor) -> ValueKind.NON_INLINEABLE_ARGUMENT_FOR_INLINE_PARAMETER_CALLED_IN_SUSPEND
+                isCallSiteIsSuspend(valueParameterDescriptor) && parameterType.kotlinType?.isSuspendFunctionTypeOrSubtype == true ->
+                    ValueKind.NON_INLINEABLE_ARGUMENT_FOR_INLINE_PARAMETER_CALLED_IN_SUSPEND
                 isInlineSuspendParameter(valueParameterDescriptor) -> ValueKind.NON_INLINEABLE_ARGUMENT_FOR_INLINE_SUSPEND_PARAMETER
                 else -> ValueKind.GENERAL
             }

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/coroutines/CoroutineTransformer.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.codegen.inline.coroutines
 
 import com.intellij.util.ArrayUtil
-import org.jetbrains.kotlin.backend.common.CodegenUtil
 import org.jetbrains.kotlin.codegen.AsmUtil.CAPTURED_THIS_FIELD
 import org.jetbrains.kotlin.codegen.ClassBuilder
 import org.jetbrains.kotlin.codegen.TransformationMethodVisitor
@@ -19,18 +18,7 @@ import org.jetbrains.kotlin.codegen.optimization.fixStack.FixStackMethodTransfor
 import org.jetbrains.kotlin.config.isReleaseCoroutines
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
-import org.jetbrains.kotlin.load.java.JvmAnnotationNames
-import org.jetbrains.kotlin.load.kotlin.FileBasedKotlinClass
-import org.jetbrains.kotlin.load.kotlin.header.KotlinClassHeader
-import org.jetbrains.kotlin.load.kotlin.header.ReadKotlinClassHeaderAnnotationVisitor
-import org.jetbrains.kotlin.metadata.ProtoBuf
-import org.jetbrains.kotlin.metadata.deserialization.*
-import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmProtoBufUtil
-import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin
-import org.jetbrains.kotlin.serialization.deserialization.getClassId
-import org.jetbrains.kotlin.serialization.deserialization.getName
 import org.jetbrains.kotlin.utils.addToStdlib.cast
 import org.jetbrains.org.objectweb.asm.*
 import org.jetbrains.org.objectweb.asm.tree.*
@@ -291,8 +279,7 @@ class SurroundSuspendLambdaCallsWithSuspendMarkersMethodVisitor(
 }
 
 private fun FunctionalArgument.isSuspendLambda(): Boolean =
-    (this is NonInlineableArgumentForInlineableParameterCalledInSuspend && isSuspend) ||
+    this is NonInlineableArgumentForInlineableParameterCalledInSuspend || (this is ExpressionLambda && isSuspend)
-            (this is PsiExpressionLambda && isSuspend)
 
 fun surroundInvokesWithSuspendMarkers(
     methodNode: MethodNode,

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.backend.jvm.codegen
 
 import org.jetbrains.kotlin.backend.common.CodegenUtil
+import org.jetbrains.kotlin.backend.common.lower.LocalDeclarationsLowering
 import org.jetbrains.kotlin.backend.common.lower.allOverridden
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
@@ -17,11 +18,11 @@ import org.jetbrains.kotlin.codegen.coroutines.reportSuspensionPointInsideMonito
 import org.jetbrains.kotlin.codegen.inline.addFakeContinuationConstructorCallMarker
 import org.jetbrains.kotlin.config.isReleaseCoroutines
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
-import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
+import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.declarations.IrFunction
-import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
-import org.jetbrains.kotlin.ir.declarations.name
 import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrGetField
+import org.jetbrains.kotlin.ir.expressions.IrGetValue
+import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
 import org.jetbrains.kotlin.ir.expressions.impl.IrErrorExpressionImpl
 import org.jetbrains.kotlin.ir.types.createType
 import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
@@ -135,6 +136,15 @@ internal fun IrFunction.isKnownToBeTailCall(): Boolean =
 internal fun IrFunction.shouldNotContainSuspendMarkers(): Boolean =
     isInvokeSuspendOfContinuation() || isKnownToBeTailCall()
 
+internal fun IrExpression?.isReadOfCrossinline(): Boolean = when (this) {
+    is IrGetValue -> (symbol.owner as? IrValueParameter)?.isCrossinline == true
+    is IrGetField -> symbol.owner.origin == LocalDeclarationsLowering.DECLARATION_ORIGIN_FIELD_FOR_CROSSINLINE_CAPTURED_VALUE
+    else -> false
+}
+
+internal fun IrExpression?.isReadOfInlineLambda(): Boolean = isReadOfCrossinline() ||
+        (this is IrGetValue && origin == IrStatementOrigin.VARIABLE_AS_FUNCTION && (symbol.owner as? IrValueParameter)?.isNoinline == false)
+
 internal fun createFakeContinuation(context: JvmBackendContext): IrExpression = IrErrorExpressionImpl(
     UNDEFINED_OFFSET,
     UNDEFINED_OFFSET,

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

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.backend.jvm.codegen
 
 import org.jetbrains.kotlin.backend.common.lower.BOUND_RECEIVER_PARAMETER
-import org.jetbrains.kotlin.backend.common.lower.LocalDeclarationsLowering
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin.*
 import org.jetbrains.kotlin.backend.jvm.intrinsics.IrIntrinsicMethods
@@ -476,21 +475,15 @@ class ExpressionCodegen(
         }
     }
 
-    private fun IrFunctionAccessExpression.isInvokeOfInlineLambda(): Boolean = when (val receiver = dispatchReceiver) {
-        is IrGetField -> receiver.symbol.owner.origin == LocalDeclarationsLowering.DECLARATION_ORIGIN_FIELD_FOR_CROSSINLINE_CAPTURED_VALUE
-        is IrGetValue -> receiver.origin == IrStatementOrigin.VARIABLE_AS_FUNCTION &&
-                (receiver.symbol.owner as? IrValueParameter)?.isNoinline == false
-        else -> false
-    }
-
     private fun IrFunctionAccessExpression.isSuspensionPoint(): Boolean = when {
         !symbol.owner.isSuspend || irFunction.shouldNotContainSuspendMarkers() -> false
         // Copy-pasted bytecode blocks are not suspension points.
         symbol.owner.isInline ->
             symbol.owner.fqNameForIrSerialization == FqName("kotlin.coroutines.intrinsics.IntrinsicsKt.suspendCoroutineUninterceptedOrReturn")
         // This includes inline lambdas, but only in functions intended for the inliner; in others, they stay as `f.invoke()`.
-        isInvokeOfInlineLambda() ->
+        dispatchReceiver.isReadOfInlineLambda() ->
-            irFunction.origin != FOR_INLINE_STATE_MACHINE_TEMPLATE && irFunction.origin != FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
+            irFunction.origin != FOR_INLINE_STATE_MACHINE_TEMPLATE && irFunction.origin != FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE &&
+                    irFunction.origin != IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA
         else -> true
     }
 

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

@@ -141,14 +141,13 @@ open class FunctionCodegen(
             // This is suspend lambda parameter of inline function
             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
-            } &&
+            }
-            // $$forInline functions never have a continuation
-            origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE
 
     private fun continuationClass() =
         irFunction.body!!.statements.firstIsInstance<IrClass>()

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

@@ -241,11 +241,13 @@ class IrExpressionLambdaImpl(
     override val lambdaClassType: Type =
         context.getLocalClassType(reference) ?: throw AssertionError("callable reference ${reference.dump()} has no name in context")
 
-    override val capturedVars: List<CapturedParamDesc> =
+    private val capturedParameters: Map<CapturedParamDesc, IrValueParameter> =
-        reference.getArgumentsWithIr().map { (param, _) ->
+        reference.getArgumentsWithIr().associate { (param, _) ->
-            capturedParamDesc(param.name.asString(), typeMapper.mapType(param.type))
+            capturedParamDesc(param.name.asString(), typeMapper.mapType(param.type)) to param
         }
 
+    override val capturedVars: List<CapturedParamDesc> = capturedParameters.keys.toList()
+
     private val loweredMethod = methodSignatureMapper.mapAsmMethod(function.suspendFunctionViewOrStub(context))
 
     val capturedParamsInDesc: List<Type> = if (isBoundCallableReference) {
@@ -271,6 +273,9 @@ class IrExpressionLambdaImpl(
     override fun getInlineSuspendLambdaViewDescriptor(): FunctionDescriptor {
         return function.suspendFunctionViewOrStub(context).descriptor
     }
+
+    override fun isCapturedSuspend(desc: CapturedParamDesc, inliningContext: InliningContext): Boolean =
+        capturedParameters[desc]?.let { it.isInlineParameter() && it.type.isSuspendFunctionTypeOrSubtype() } == true
 }
 
 class IrDefaultLambda(

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

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
 import org.jetbrains.kotlin.backend.jvm.codegen.isInvokeSuspendForInlineOfLambda
 import org.jetbrains.kotlin.backend.jvm.codegen.isInvokeSuspendOfContinuation
 import org.jetbrains.kotlin.backend.jvm.codegen.isInvokeSuspendOfLambda
+import org.jetbrains.kotlin.backend.jvm.codegen.isReadOfCrossinline
 import org.jetbrains.kotlin.backend.jvm.ir.*
 import org.jetbrains.kotlin.backend.jvm.localDeclarationsPhase
 import org.jetbrains.kotlin.builtins.functions.FunctionInvokeDescriptor.Factory.BIG_ARITY
@@ -679,13 +680,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
     private class SuspendLambdaInfo(val reference: IrFunctionReference) {
         val function = reference.symbol.owner
         val arity = (reference.type as IrSimpleType).arguments.size - 1
-        val capturesCrossinline = function.valueParameters.any {
+        val capturesCrossinline = function.valueParameters.any { reference.getValueArgument(it.index).isReadOfCrossinline() }
-            when (val argument = reference.getValueArgument(it.index)) {
-                is IrGetValue -> (argument.symbol.owner as? IrValueParameter)?.isCrossinline == true
-                is IrGetField -> argument.symbol.owner.origin == LocalDeclarationsLowering.DECLARATION_ORIGIN_FIELD_FOR_CROSSINLINE_CAPTURED_VALUE
-                else -> false
-            }
-        }
         lateinit var constructor: IrConstructor
     }
 }

compiler/testData/codegen/box/coroutines/tailCallOptimizations/crossinline_ir.txt

@@ -44,20 +44,6 @@ public final class CrossinlineKt$box$1$invokeSuspend$$inlined$filter$1 {
     public @org.jetbrains.annotations.Nullable method consume(@org.jetbrains.annotations.NotNull p0: Sink, @org.jetbrains.annotations.NotNull p1: kotlin.coroutines.Continuation): java.lang.Object
 }
 
-@kotlin.Metadata
-@kotlin.coroutines.jvm.internal.DebugMetadata
-final class CrossinlineKt$box$1$invokeSuspend$$inlined$fold$1$1 {
-    field L$0: java.lang.Object
-    field L$1: java.lang.Object
-    field L$2: java.lang.Object
-    field L$3: java.lang.Object
-    field label: int
-    @org.jetbrains.annotations.NotNull field result: java.lang.Object
-    synthetic final field this$0: CrossinlineKt$box$1$invokeSuspend$$inlined$fold$1
-    public method <init>(p0: CrossinlineKt$box$1$invokeSuspend$$inlined$fold$1, p1: kotlin.coroutines.Continuation): void
-    public final @org.jetbrains.annotations.Nullable method invokeSuspend(@org.jetbrains.annotations.NotNull p0: java.lang.Object): java.lang.Object
-}
-
 @kotlin.Metadata
 public final class CrossinlineKt$box$1$invokeSuspend$$inlined$fold$1 {
     synthetic final field $acc$inlined: kotlin.jvm.internal.Ref$ObjectRef

compiler/testData/codegen/boxInline/suspend/crossinlineSuspendLambdaInsideCrossinlineSuspendLambda.kt

@@ -2,6 +2,8 @@
 // COMMON_COROUTINES_TEST
 // WITH_RUNTIME
 // WITH_COROUTINES
+// IGNORE_BACKEND: JVM_IR
+// IGNORE_BACKEND_MULTI_MODULE: JVM_IR
 
 import COROUTINES_PACKAGE.*
 import helpers.*

compiler/testData/codegen/boxInline/suspend/doubleRegenerationWithNonSuspendingLambda.kt

@@ -2,8 +2,6 @@
 // COMMON_COROUTINES_TEST
 // WITH_RUNTIME
 // WITH_COROUTINES
-// IGNORE_BACKEND: JVM_IR
-// IGNORE_BACKEND_MULTI_MODULE: JVM_IR
 
 interface Flow<T> {
     abstract suspend fun collect(collector: FlowCollector<T>)