JVM_IR: do not place suspend markers around crossinline lambda calls

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.IrFunction
-import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
-import org.jetbrains.kotlin.ir.declarations.name
+import org.jetbrains.kotlin.ir.declarations.*
 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() ->
-            irFunction.origin != FOR_INLINE_STATE_MACHINE_TEMPLATE && irFunction.origin != FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
+        dispatchReceiver.isReadOfInlineLambda() ->
+            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> =
-        reference.getArgumentsWithIr().map { (param, _) ->
-            capturedParamDesc(param.name.asString(), typeMapper.mapType(param.type))
+    private val capturedParameters: Map<CapturedParamDesc, IrValueParameter> =
+        reference.getArgumentsWithIr().associate { (param, _) ->
+            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 {
-            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
-            }
-        }
+        val capturesCrossinline = function.valueParameters.any { reference.getValueArgument(it.index).isReadOfCrossinline() }
         lateinit var constructor: IrConstructor
     }
 }