JVM_IR: do not place suspend markers around crossinline lambda calls

2020-02-20 11:48:34 +01:00
parent 39ebc8cfa5
commit 6b98ea2378
13 changed files with 43 additions and 65 deletions
@@ -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 }) {
-                        recapturedParamInfo.functionalArgument = NonInlineableArgumentForInlineableParameterCalledInSuspend(
-                            isCapturedSuspendLambda(info.closure, desc.fieldName, inliningContext.state.bindingContext)
-                        )
+                    if (info is ExpressionLambda && info.isCapturedSuspend(desc, inliningContext)) {
+                        recapturedParamInfo.functionalArgument = NonInlineableArgumentForInlineableParameterCalledInSuspend
                    }
                    val composed = StackValue.field(
                        desc.type,
@@ -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
                }
@@ -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)
 }
@@ -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
            }
@@ -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 PsiExpressionLambda && isSuspend)
+    this is NonInlineableArgumentForInlineableParameterCalledInSuspend || (this is ExpressionLambda && isSuspend)

 fun surroundInvokesWithSuspendMarkers(
    methodNode: MethodNode,
@@ -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,
@@ -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
    }

@@ -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>()
@@ -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(
@@ -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
    }
 }
@@ -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
@@ -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.*
@@ -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>)