JVM: use conditional suspension point markers

2020-04-02 13:03:12 +02:00
parent 42a48ef312
commit c650c9570f
18 changed files with 125 additions and 333 deletions
@@ -28,6 +28,7 @@ import org.jetbrains.kotlin.codegen.context.*;
 import org.jetbrains.kotlin.codegen.coroutines.CoroutineCodegenForLambda;
 import org.jetbrains.kotlin.codegen.coroutines.CoroutineCodegenUtilKt;
 import org.jetbrains.kotlin.codegen.coroutines.ResolvedCallWithRealDescriptor;
+import org.jetbrains.kotlin.codegen.coroutines.SuspensionPointKind;
 import org.jetbrains.kotlin.codegen.extensions.ExpressionCodegenExtension;
 import org.jetbrains.kotlin.codegen.inline.*;
 import org.jetbrains.kotlin.codegen.intrinsics.*;
@@ -2592,11 +2593,12 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
            return;
        }

-        boolean isSuspendNoInlineCall =
-                CoroutineCodegenUtilKt.isSuspendNoInlineCall(resolvedCall, this, state.getLanguageVersionSettings());
+        SuspensionPointKind suspensionPointKind =
+                CoroutineCodegenUtilKt.isSuspensionPoint(resolvedCall, this, state.getLanguageVersionSettings());
+        boolean maybeSuspensionPoint = suspensionPointKind != SuspensionPointKind.NEVER;
        boolean isConstructor = resolvedCall.getResultingDescriptor() instanceof ConstructorDescriptor;
        if (!(callableMethod instanceof IntrinsicWithSpecialReceiver)) {
-            putReceiverAndInlineMarkerIfNeeded(callableMethod, resolvedCall, receiver, isSuspendNoInlineCall, isConstructor);
+            putReceiverAndInlineMarkerIfNeeded(callableMethod, resolvedCall, receiver, maybeSuspensionPoint, isConstructor);
        }

        callGenerator.processAndPutHiddenParameters(false);
@@ -2630,16 +2632,16 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
            }
        }

-        if (isSuspendNoInlineCall) {
-            addSuspendMarker(v, true);
+        if (maybeSuspensionPoint) {
+            addSuspendMarker(v, suspensionPointKind, true);
        }

        callGenerator.genCall(callableMethod, resolvedCall, defaultMaskWasGenerated, this);

-        if (isSuspendNoInlineCall) {
+        if (maybeSuspensionPoint) {
            addReturnsUnitMarkerIfNecessary(v, resolvedCall);

-            addSuspendMarker(v, false);
+            addSuspendMarker(v, suspensionPointKind, false);
            addInlineMarker(v, false);
        }

@@ -19,7 +19,6 @@ import org.jetbrains.kotlin.codegen.binding.CodegenBinding;
 import org.jetbrains.kotlin.codegen.context.*;
 import org.jetbrains.kotlin.codegen.coroutines.CoroutineCodegenUtilKt;
 import org.jetbrains.kotlin.codegen.coroutines.SuspendFunctionGenerationStrategy;
-import org.jetbrains.kotlin.codegen.coroutines.SuspendInlineFunctionGenerationStrategy;
 import org.jetbrains.kotlin.codegen.state.GenerationState;
 import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper;
 import org.jetbrains.kotlin.codegen.state.TypeMapperUtilsKt;
@@ -127,17 +126,8 @@ public class FunctionCodegen {
            if (functionDescriptor.isSuspend()) {
                if (isEffectivelyInlineOnly(functionDescriptor)) {
                    strategy = new FunctionGenerationStrategy.FunctionDefault(state, function);
-                } else if (!functionDescriptor.isInline()) {
-                    strategy = new SuspendFunctionGenerationStrategy(
-                            state,
-                            CoroutineCodegenUtilKt.<FunctionDescriptor>unwrapInitialDescriptorForSuspendFunction(functionDescriptor),
-                            function,
-                            v.getThisName(),
-                            state.getConstructorCallNormalizationMode(),
-                            this
-                    );
                } else {
-                    strategy = new SuspendInlineFunctionGenerationStrategy(
+                    strategy = new SuspendFunctionGenerationStrategy(
                            state,
                            CoroutineCodegenUtilKt.<FunctionDescriptor>unwrapInitialDescriptorForSuspendFunction(functionDescriptor),
                            function,
@@ -15,7 +15,6 @@ import org.jetbrains.kotlin.codegen.binding.CodegenBinding.CAPTURES_CROSSINLINE_
 import org.jetbrains.kotlin.codegen.binding.CodegenBinding.CLOSURE
 import org.jetbrains.kotlin.codegen.context.ClosureContext
 import org.jetbrains.kotlin.codegen.context.MethodContext
-import org.jetbrains.kotlin.codegen.inline.coroutines.SurroundSuspendLambdaCallsWithSuspendMarkersMethodVisitor
 import org.jetbrains.kotlin.codegen.serialization.JvmSerializationBindings.METHOD_FOR_FUNCTION
 import org.jetbrains.kotlin.codegen.serialization.JvmSerializerExtension
 import org.jetbrains.kotlin.codegen.state.GenerationState
@@ -501,17 +500,11 @@ class CoroutineCodegenForLambda private constructor(
                        languageVersionSettings = languageVersionSettings,
                        disableTailCallOptimizationForFunctionReturningUnit = false
                    )
-                    return if (forInline) AddEndLabelMethodVisitor(
-                        MethodNodeCopyingMethodVisitor(
-                            SurroundSuspendLambdaCallsWithSuspendMarkersMethodVisitor(
-                                stateMachineBuilder, access, name, desc, v.thisName,
-                                isCapturedSuspendLambda = { isCapturedSuspendLambda(closure, it.name, state.bindingContext) }
-                            ), access, name, desc,
-                            newMethod = { origin, newAccess, newName, newDesc ->
-                                functionCodegen.newMethod(origin, newAccess, newName, newDesc, null, null)
-                            }
-                        ), access, name, desc, endLabel
-                    ) else AddEndLabelMethodVisitor(stateMachineBuilder, access, name, desc, endLabel)
+                    val maybeWithForInline = if (forInline)
+                        SuspendForInlineCopyingMethodVisitor(stateMachineBuilder, access, name, desc, functionCodegen::newMethod)
+                    else
+                        stateMachineBuilder
+                    return AddEndLabelMethodVisitor(maybeWithForInline, access, name, desc, endLabel)
                }

                override fun doGenerateBody(codegen: ExpressionCodegen, signature: JvmMethodSignature) {
@@ -9,7 +9,8 @@ import org.jetbrains.kotlin.backend.common.CodegenUtil
 import org.jetbrains.kotlin.codegen.*
 import org.jetbrains.kotlin.codegen.binding.CodegenBinding
 import org.jetbrains.kotlin.codegen.inline.addFakeContinuationConstructorCallMarker
-import org.jetbrains.kotlin.codegen.inline.coroutines.SurroundSuspendLambdaCallsWithSuspendMarkersMethodVisitor
+import org.jetbrains.kotlin.codegen.inline.coroutines.FOR_INLINE_SUFFIX
+import org.jetbrains.kotlin.codegen.inline.preprocessSuspendMarkers
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.config.JVMConstructorCallNormalizationMode
 import org.jetbrains.kotlin.config.LanguageVersionSettings
@@ -18,23 +19,23 @@ import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.psi.KtFunction
 import org.jetbrains.kotlin.psi.psiUtil.getElementTextWithContext
+import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.OtherOrigin
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature
 import org.jetbrains.kotlin.types.typeUtil.isUnit
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
-import org.jetbrains.kotlin.utils.sure
 import org.jetbrains.org.objectweb.asm.MethodVisitor
 import org.jetbrains.org.objectweb.asm.Opcodes
 import org.jetbrains.org.objectweb.asm.Type
 import org.jetbrains.org.objectweb.asm.tree.MethodNode

-open class SuspendFunctionGenerationStrategy(
+class SuspendFunctionGenerationStrategy(
    state: GenerationState,
-    protected val originalSuspendDescriptor: FunctionDescriptor,
-    protected val declaration: KtFunction,
-    protected val containingClassInternalName: String,
+    private val originalSuspendDescriptor: FunctionDescriptor,
+    private val declaration: KtFunction,
+    private val containingClassInternalName: String,
    private val constructorCallNormalizationMode: JVMConstructorCallNormalizationMode,
-    protected val functionCodegen: FunctionCodegen
+    private val functionCodegen: FunctionCodegen
 ) : FunctionGenerationStrategy.CodegenBased(state) {

    private lateinit var codegen: ExpressionCodegen
@@ -56,39 +57,23 @@ open class SuspendFunctionGenerationStrategy(
        if (access and Opcodes.ACC_ABSTRACT != 0) return mv

        val stateMachineBuilder = createStateMachineBuilder(mv, access, name, desc)
-
-        val forInline = state.bindingContext[CodegenBinding.CAPTURES_CROSSINLINE_LAMBDA, originalSuspendDescriptor] == true
-        // Both capturing and inline functions share the same suffix, however, inline functions can also be capturing
-        // they are already covered by SuspendInlineFunctionGenerationStrategy, thus, if we generate yet another copy,
-        // we will get name+descriptor clash
-        return if (forInline && !originalSuspendDescriptor.isInline)
-            AddConstructorCallForCoroutineRegeneration(
-                MethodNodeCopyingMethodVisitor(
-                    SurroundSuspendLambdaCallsWithSuspendMarkersMethodVisitor(
-                        stateMachineBuilder,
-                        access, name, desc, containingClassInternalName,
-                        isCapturedSuspendLambda = {
-                            isCapturedSuspendLambda(
-                                functionCodegen.closure.sure {
-                                    "Anonymous object should have closure"
-                                },
-                                it.name,
-                                state.bindingContext
-                            )
-                        }
-                    ), access, name, desc,
-                    newMethod = { origin, newAccess, newName, newDesc ->
-                        functionCodegen.newMethod(origin, newAccess, newName, newDesc, null, null)
-                    }
-                ), access, name, desc, null, null, this::classBuilderForCoroutineState,
+        if (originalSuspendDescriptor.isInline) {
+            return SuspendForInlineCopyingMethodVisitor(stateMachineBuilder, access, name, desc, functionCodegen::newMethod, keepAccess = false)
+        }
+        if (state.bindingContext[CodegenBinding.CAPTURES_CROSSINLINE_LAMBDA, originalSuspendDescriptor] == true) {
+            return AddConstructorCallForCoroutineRegeneration(
+                SuspendForInlineCopyingMethodVisitor(stateMachineBuilder, access, name, desc, functionCodegen::newMethod),
+                access, name, desc, null, null, this::classBuilderForCoroutineState,
                containingClassInternalName,
                originalSuspendDescriptor.dispatchReceiverParameter != null,
                containingClassInternalNameOrNull(),
                languageVersionSettings
-            ) else stateMachineBuilder
+            )
+        }
+        return stateMachineBuilder
    }

-    protected fun createStateMachineBuilder(
+    private fun createStateMachineBuilder(
        mv: MethodVisitor,
        access: Int,
        name: String,
@@ -170,3 +155,28 @@ open class SuspendFunctionGenerationStrategy(
        }
    }
 }
+
+// For named suspend function we generate two methods:
+// 1) to use as noinline function, which have state machine
+// 2) to use from inliner: private one without state machine
+class SuspendForInlineCopyingMethodVisitor(
+    delegate: MethodVisitor, access: Int, name: String, desc: String,
+    private val newMethod: (JvmDeclarationOrigin, Int, String, String, String?, Array<String>?) -> MethodVisitor,
+    private val keepAccess: Boolean = true
+) : TransformationMethodVisitor(delegate, access, name, desc, null, null) {
+    override fun performTransformations(methodNode: MethodNode) {
+        val newMethodNode = with(methodNode) {
+            val newAccess = if (keepAccess) access else
+                access or Opcodes.ACC_PRIVATE and Opcodes.ACC_PUBLIC.inv() and Opcodes.ACC_PROTECTED.inv()
+            MethodNode(newAccess, name + FOR_INLINE_SUFFIX, desc, signature, exceptions.toTypedArray())
+        }
+        val newMethodVisitor = with(newMethodNode) {
+            newMethod(JvmDeclarationOrigin.NO_ORIGIN, access, name, desc, signature, exceptions.toTypedArray())
+        }
+        methodNode.instructions.resetLabels()
+        methodNode.accept(newMethodNode)
+        methodNode.preprocessSuspendMarkers(forInline = false, keepFakeContinuation = false)
+        newMethodNode.preprocessSuspendMarkers(forInline = true, keepFakeContinuation = true)
+        newMethodNode.accept(newMethodVisitor)
+    }
+}
@@ -1,129 +0,0 @@
-/*
- * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
- * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
- */
-
-package org.jetbrains.kotlin.codegen.coroutines
-
-import org.jetbrains.kotlin.builtins.isSuspendFunctionType
-import org.jetbrains.kotlin.codegen.ExpressionCodegen
-import org.jetbrains.kotlin.codegen.FunctionCodegen
-import org.jetbrains.kotlin.codegen.FunctionGenerationStrategy
-import org.jetbrains.kotlin.codegen.TransformationMethodVisitor
-import org.jetbrains.kotlin.codegen.inline.coroutines.FOR_INLINE_SUFFIX
-import org.jetbrains.kotlin.codegen.inline.coroutines.findReceiverOfInvoke
-import org.jetbrains.kotlin.codegen.inline.coroutines.surroundInvokesWithSuspendMarkers
-import org.jetbrains.kotlin.codegen.inline.isInvokeOnLambda
-import org.jetbrains.kotlin.codegen.optimization.common.asSequence
-import org.jetbrains.kotlin.codegen.optimization.fixStack.FixStackMethodTransformer
-import org.jetbrains.kotlin.codegen.optimization.transformer.MethodTransformer
-import org.jetbrains.kotlin.codegen.state.GenerationState
-import org.jetbrains.kotlin.config.JVMConstructorCallNormalizationMode
-import org.jetbrains.kotlin.descriptors.FunctionDescriptor
-import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
-import org.jetbrains.kotlin.psi.KtFunction
-import org.jetbrains.kotlin.resolve.inline.InlineUtil
-import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin
-import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature
-import org.jetbrains.org.objectweb.asm.MethodVisitor
-import org.jetbrains.org.objectweb.asm.Opcodes
-import org.jetbrains.org.objectweb.asm.tree.*
-import org.jetbrains.org.objectweb.asm.tree.analysis.SourceInterpreter
-
-// For named suspend function we generate two methods:
-// 1) to use as noinline function, which have state machine
-// 2) to use from inliner: private one without state machine
-class SuspendInlineFunctionGenerationStrategy(
-    state: GenerationState,
-    originalSuspendDescriptor: FunctionDescriptor,
-    declaration: KtFunction,
-    containingClassInternalName: String,
-    constructorCallNormalizationMode: JVMConstructorCallNormalizationMode,
-    codegen: FunctionCodegen
-) : SuspendFunctionGenerationStrategy(
-    state,
-    originalSuspendDescriptor,
-    declaration,
-    containingClassInternalName,
-    constructorCallNormalizationMode,
-    codegen
-) {
-    private val defaultStrategy = FunctionGenerationStrategy.FunctionDefault(state, declaration)
-
-    override fun wrapMethodVisitor(mv: MethodVisitor, access: Int, name: String, desc: String): MethodVisitor {
-        if (access and Opcodes.ACC_ABSTRACT != 0) return mv
-
-        return MethodNodeCopyingMethodVisitor(
-            SurroundSuspendParameterCallsWithSuspendMarkersMethodVisitor(
-                createStateMachineBuilder(mv, access, name, desc),
-                access, name, desc, containingClassInternalName, originalSuspendDescriptor.valueParameters
-            ), access, name, desc,
-            newMethod = { origin, newAccess, newName, newDesc ->
-                functionCodegen.newMethod(origin, newAccess, newName, newDesc, null, null)
-            }, keepAccess = false
-        )
-    }
-
-    override fun doGenerateBody(codegen: ExpressionCodegen, signature: JvmMethodSignature) {
-        super.doGenerateBody(codegen, signature)
-        defaultStrategy.doGenerateBody(codegen, signature)
-    }
-}
-
-class MethodNodeCopyingMethodVisitor(
-    delegate: MethodVisitor, private val access: Int, private val name: String, private val desc: String,
-    private val newMethod: (JvmDeclarationOrigin, Int, String, String) -> MethodVisitor,
-    private val keepAccess: Boolean = true
-) : TransformationMethodVisitor(
-    delegate, calculateAccessForInline(access, keepAccess), "$name$FOR_INLINE_SUFFIX", desc, null, null
-) {
-    override fun performTransformations(methodNode: MethodNode) {
-        val newMethodNode = newMethod(
-            JvmDeclarationOrigin.NO_ORIGIN, calculateAccessForInline(access, keepAccess), "$name$FOR_INLINE_SUFFIX", desc
-        )
-        methodNode.instructions.resetLabels()
-        methodNode.accept(newMethodNode)
-    }
-
-    companion object {
-        private fun calculateAccessForInline(access: Int, keepAccess: Boolean): Int =
-            if (keepAccess) access
-            else access or Opcodes.ACC_PRIVATE and Opcodes.ACC_PUBLIC.inv() and Opcodes.ACC_PROTECTED.inv()
-    }
-}
-
-private class SurroundSuspendParameterCallsWithSuspendMarkersMethodVisitor(
-    delegate: MethodVisitor,
-    access: Int,
-    name: String,
-    desc: String,
-    private val thisName: String,
-    private val valueParameters: List<ValueParameterDescriptor>
-): TransformationMethodVisitor(delegate, access, name, desc, null, null) {
-    override fun performTransformations(methodNode: MethodNode) {
-        fun AbstractInsnNode.index() = methodNode.instructions.indexOf(this)
-        fun AbstractInsnNode.isInlineSuspendParameter(): Boolean {
-            if (this !is VarInsnNode) return false
-            val index = `var` - (if (methodNode.access and Opcodes.ACC_STATIC != 0) 0 else 1)
-            return opcode == Opcodes.ALOAD && index < valueParameters.size && InlineUtil.isInlineParameter(valueParameters[index]) &&
-                    valueParameters[index].type.isSuspendFunctionType
-        }
-
-        FixStackMethodTransformer().transform(thisName, methodNode)
-
-        val sourceFrames = MethodTransformer.analyze(thisName, methodNode, SourceInterpreter())
-
-        val noinlineInvokes = arrayListOf<Pair<AbstractInsnNode, AbstractInsnNode>>()
-
-        for (insn in methodNode.instructions.asSequence()) {
-            if (insn.opcode != Opcodes.INVOKEINTERFACE) continue
-            insn as MethodInsnNode
-            if (!isInvokeOnLambda(insn.owner, insn.name)) continue
-            val frame = sourceFrames[insn.index()] ?: continue
-            val aload = findReceiverOfInvoke(frame, insn).takeIf { it?.isInlineSuspendParameter() == true } as? VarInsnNode ?: continue
-            noinlineInvokes.add(insn to aload)
-        }
-
-        surroundInvokesWithSuspendMarkers(methodNode, noinlineInvokes)
-    }
-}
@@ -219,15 +219,18 @@ private fun NewResolvedCallImpl<VariableDescriptor>.asDummyOldResolvedCall(bindi
    )
 }

-fun ResolvedCall<*>.isSuspendNoInlineCall(codegen: ExpressionCodegen, languageVersionSettings: LanguageVersionSettings): Boolean {
+enum class SuspensionPointKind { NEVER, NOT_INLINE, ALWAYS }
+
+fun ResolvedCall<*>.isSuspensionPoint(codegen: ExpressionCodegen, languageVersionSettings: LanguageVersionSettings): SuspensionPointKind {
+    val functionDescriptor = resultingDescriptor as? FunctionDescriptor ?: return SuspensionPointKind.NEVER
+    if (!functionDescriptor.unwrapInitialDescriptorForSuspendFunction().isSuspend) return SuspensionPointKind.NEVER
+    if (functionDescriptor.isBuiltInSuspendCoroutineUninterceptedOrReturnInJvm(languageVersionSettings)) return SuspensionPointKind.ALWAYS
+    if (functionDescriptor.isInline) return SuspensionPointKind.NEVER
+
    val isInlineLambda = this.safeAs<VariableAsFunctionResolvedCall>()
        ?.variableCall?.resultingDescriptor?.safeAs<ValueParameterDescriptor>()
        ?.let { it.isCrossinline || (!it.isNoinline && codegen.context.functionDescriptor.isInline) } == true
-
-    val functionDescriptor = resultingDescriptor as? FunctionDescriptor ?: return false
-    if (!functionDescriptor.unwrapInitialDescriptorForSuspendFunction().isSuspend) return false
-    if (functionDescriptor.isBuiltInSuspendCoroutineUninterceptedOrReturnInJvm(languageVersionSettings)) return true
-    return !(functionDescriptor.isInline || isInlineLambda)
+    return if (isInlineLambda) SuspensionPointKind.NOT_INLINE else SuspensionPointKind.ALWAYS
 }

 fun CallableDescriptor.isSuspendFunctionNotSuspensionView(): Boolean {
@@ -534,7 +534,9 @@ abstract class InlineCodegen<out T : BaseExpressionCodegen>(

            val directMember = getDirectMemberAndCallableFromObject(functionDescriptor)
            if (!isBuiltInArrayIntrinsic(functionDescriptor) && directMember !is DescriptorWithContainerSource) {
-                return sourceCompilerForInline.doCreateMethodNodeFromSource(functionDescriptor, jvmSignature, callDefault, asmMethod)
+                val node = sourceCompilerForInline.doCreateMethodNodeFromSource(functionDescriptor, jvmSignature, callDefault, asmMethod)
+                node.node.preprocessSuspendMarkers(forInline = true, keepFakeContinuation = false)
+                return node
            }

            return getCompiledMethodNodeInner(functionDescriptor, directMember, asmMethod, methodOwner, state, jvmSignature)
@@ -222,6 +222,7 @@ internal fun Type.boxReceiverForBoundReference(kotlinType: KotlinType, typeMappe
 abstract class ExpressionLambda(isCrossInline: Boolean) : LambdaInfo(isCrossInline) {
    override fun generateLambdaBody(sourceCompiler: SourceCompilerForInline, reifiedTypeInliner: ReifiedTypeInliner<*>) {
        node = sourceCompiler.generateLambdaBody(this)
+        node.node.preprocessSuspendMarkers(forInline = true, keepFakeContinuation = false)
    }

    abstract fun getInlineSuspendLambdaViewDescriptor(): FunctionDescriptor
@@ -107,15 +107,7 @@ class CoroutineTransformer(
            )

            if (generateForInline)
-                MethodNodeCopyingMethodVisitor(
-                    delegate = stateMachineBuilder,
-                    access = node.access,
-                    name = name,
-                    desc = node.desc,
-                    newMethod = { origin, newAccess, newName, newDesc ->
-                        classBuilder.newMethod(origin, newAccess, newName, newDesc, null, null)
-                    }
-                )
+                SuspendForInlineCopyingMethodVisitor(stateMachineBuilder, node.access, name, node.desc, classBuilder::newMethod)
            else
                stateMachineBuilder
        }
@@ -154,12 +146,7 @@ class CoroutineTransformer(
            )

            if (generateForInline)
-                MethodNodeCopyingMethodVisitor(
-                    stateMachineBuilder, node.access, name, node.desc,
-                    newMethod = { origin, newAccess, newName, newDesc ->
-                        classBuilder.newMethod(origin, newAccess, newName, newDesc, null, null)
-                    }
-                )
+                SuspendForInlineCopyingMethodVisitor(stateMachineBuilder, node.access, name, node.desc, classBuilder::newMethod)
            else
                stateMachineBuilder
        }
@@ -18,7 +18,10 @@ 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.SuspensionPointKind
 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
@@ -71,8 +74,8 @@ private const val INLINE_MARKER_AFTER_METHOD_NAME = "afterInlineCall"
 private const val INLINE_MARKER_FINALLY_START = "finallyStart"

 private const val INLINE_MARKER_FINALLY_END = "finallyEnd"
-const val INLINE_MARKER_BEFORE_SUSPEND_ID = 0
-const val INLINE_MARKER_AFTER_SUSPEND_ID = 1
+private const val INLINE_MARKER_BEFORE_SUSPEND_ID = 0
+private const val INLINE_MARKER_AFTER_SUSPEND_ID = 1
 private const val INLINE_MARKER_RETURNS_UNIT = 2
 private const val INLINE_MARKER_FAKE_CONTINUATION = 3
 private const val INLINE_MARKER_BEFORE_FAKE_CONTINUATION_CONSTRUCTOR_CALL = 4
@@ -437,8 +440,10 @@ fun addInlineSuspendMarker(v: InstructionAdapter, isStartNotEnd: Boolean) {
    v.emitInlineMarker(if (isStartNotEnd) INLINE_MARKER_BEFORE_INLINE_SUSPEND_ID else INLINE_MARKER_AFTER_INLINE_SUSPEND_ID)
 }

-fun addSuspendMarker(v: InstructionAdapter, isStartNotEnd: Boolean, isInline: Boolean) {
-    if (isInline) addInlineSuspendMarker(v, isStartNotEnd) else addSuspendMarker(v, isStartNotEnd)
+fun addSuspendMarker(v: InstructionAdapter, kind: SuspensionPointKind, isStartNotEnd: Boolean) = when (kind) {
+    SuspensionPointKind.NEVER -> Unit
+    SuspensionPointKind.NOT_INLINE -> addInlineSuspendMarker(v, isStartNotEnd)
+    SuspensionPointKind.ALWAYS -> addSuspendMarker(v, isStartNotEnd)
 }

 fun addFakeContinuationConstructorCallMarker(v: InstructionAdapter, isStartNotEnd: Boolean) {
@@ -466,16 +471,16 @@ private fun InstructionAdapter.emitInlineMarker(id: Int) {

 internal fun isBeforeSuspendMarker(insn: AbstractInsnNode) = isSuspendMarker(insn, INLINE_MARKER_BEFORE_SUSPEND_ID)
 internal fun isAfterSuspendMarker(insn: AbstractInsnNode) = isSuspendMarker(insn, INLINE_MARKER_AFTER_SUSPEND_ID)
-fun isBeforeInlineSuspendMarker(insn: AbstractInsnNode) = isSuspendMarker(insn, INLINE_MARKER_BEFORE_INLINE_SUSPEND_ID)
-fun isAfterInlineSuspendMarker(insn: AbstractInsnNode) = isSuspendMarker(insn, INLINE_MARKER_AFTER_INLINE_SUSPEND_ID)
+internal fun isBeforeInlineSuspendMarker(insn: AbstractInsnNode) = isSuspendMarker(insn, INLINE_MARKER_BEFORE_INLINE_SUSPEND_ID)
+internal fun isAfterInlineSuspendMarker(insn: AbstractInsnNode) = isSuspendMarker(insn, INLINE_MARKER_AFTER_INLINE_SUSPEND_ID)
 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

-fun isBeforeFakeContinuationConstructorCallMarker(insn: AbstractInsnNode) =
+internal fun isBeforeFakeContinuationConstructorCallMarker(insn: AbstractInsnNode) =
    isSuspendMarker(insn, INLINE_MARKER_BEFORE_FAKE_CONTINUATION_CONSTRUCTOR_CALL)

-fun isAfterFakeContinuationConstructorCallMarker(insn: AbstractInsnNode) =
+internal fun isAfterFakeContinuationConstructorCallMarker(insn: AbstractInsnNode) =
    isSuspendMarker(insn, INLINE_MARKER_AFTER_FAKE_CONTINUATION_CONSTRUCTOR_CALL)

 private fun isSuspendMarker(insn: AbstractInsnNode, id: Int) =
@@ -563,3 +568,30 @@ fun initDefaultSourceMappingIfNeeded(
        parentContext = parentContext.parentContext
    }
 }
+
+fun MethodNode.preprocessSuspendMarkers(forInline: Boolean, keepFakeContinuation: Boolean = true) {
+    if (instructions.first == null) return
+    if (!keepFakeContinuation) {
+        val sequence = instructions.asSequence()
+        val start = sequence.find { isBeforeFakeContinuationConstructorCallMarker(it) }
+        val end = sequence.find { isAfterFakeContinuationConstructorCallMarker(it) }
+        if (start != null) {
+            // Include one instruction before the start marker (that's the id) and one after the end marker (that's a pop).
+            InsnSequence(start.previous, end?.next?.next).forEach(instructions::remove)
+        }
+    }
+    for (insn in instructions.asSequence().filter { isBeforeInlineSuspendMarker(it) || isAfterInlineSuspendMarker(it) }) {
+        if (forInline || keepFakeContinuation) {
+            val beforeMarker = insn.previous.previous
+            if (isReturnsUnitMarker(beforeMarker)) {
+                instructions.remove(beforeMarker.previous)
+                instructions.remove(beforeMarker)
+            }
+            instructions.remove(insn.previous)
+            instructions.remove(insn)
+        } else {
+            val newId = if (isBeforeInlineSuspendMarker(insn)) INLINE_MARKER_BEFORE_SUSPEND_ID else INLINE_MARKER_AFTER_SUSPEND_ID
+            instructions.set(insn.previous, InsnNode(Opcodes.ICONST_0 + newId))
+        }
+    }
+}