JVM_IR: remove $$forInline checks from ExpressionCodegen

2020-03-04 13:17:50 +01:00
parent de1fa40c7e
commit 8945d5a0df
5 changed files with 68 additions and 25 deletions
@@ -71,12 +71,14 @@ 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"
-private const val INLINE_MARKER_BEFORE_SUSPEND_ID = 0
-private const val INLINE_MARKER_AFTER_SUSPEND_ID = 1
+const val INLINE_MARKER_BEFORE_SUSPEND_ID = 0
+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
 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

 internal fun getMethodNode(
    classData: ByteArray,
@@ -431,6 +433,14 @@ fun addSuspendMarker(v: InstructionAdapter, isStartNotEnd: Boolean) {
    v.emitInlineMarker(if (isStartNotEnd) INLINE_MARKER_BEFORE_SUSPEND_ID else INLINE_MARKER_AFTER_SUSPEND_ID)
 }

+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 addFakeContinuationConstructorCallMarker(v: InstructionAdapter, isStartNotEnd: Boolean) {
    v.emitInlineMarker(if (isStartNotEnd) INLINE_MARKER_BEFORE_FAKE_CONTINUATION_CONSTRUCTOR_CALL else INLINE_MARKER_AFTER_FAKE_CONTINUATION_CONSTRUCTOR_CALL)
 }
@@ -456,14 +466,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 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 isBeforeFakeContinuationConstructorCallMarker(insn: AbstractInsnNode) =
+fun isBeforeFakeContinuationConstructorCallMarker(insn: AbstractInsnNode) =
    isSuspendMarker(insn, INLINE_MARKER_BEFORE_FAKE_CONTINUATION_CONSTRUCTOR_CALL)

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

 private fun isSuspendMarker(insn: AbstractInsnNode, id: Int) =
@@ -355,6 +355,7 @@ open class ClassCodegen protected constructor(
        }

        val node = FunctionCodegen(method, this).generate()
+        node.preprocessSuspendMarkers(method)
        val mv = with(node) { visitor.newMethod(method.OtherOrigin, access, name, desc, signature, exceptions.toTypedArray()) }
        if (method.hasContinuation() || method.isInvokeSuspendOfLambda()) {
            // Generate a state machine within this method. The continuation class for it should be generated
@@ -16,7 +16,10 @@ import org.jetbrains.kotlin.codegen.coroutines.CoroutineTransformerMethodVisitor
 import org.jetbrains.kotlin.codegen.coroutines.INVOKE_SUSPEND_METHOD_NAME
 import org.jetbrains.kotlin.codegen.coroutines.SUSPEND_IMPL_NAME_SUFFIX
 import org.jetbrains.kotlin.codegen.coroutines.reportSuspensionPointInsideMonitor
+import org.jetbrains.kotlin.codegen.inline.*
 import org.jetbrains.kotlin.codegen.inline.coroutines.FOR_INLINE_SUFFIX
+import org.jetbrains.kotlin.codegen.optimization.common.InsnSequence
+import org.jetbrains.kotlin.codegen.optimization.common.asSequence
 import org.jetbrains.kotlin.config.isReleaseCoroutines
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.declarations.*
@@ -32,6 +35,9 @@ import org.jetbrains.kotlin.ir.util.parentAsClass
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.org.objectweb.asm.MethodVisitor
+import org.jetbrains.org.objectweb.asm.Opcodes
+import org.jetbrains.org.objectweb.asm.tree.AbstractInsnNode
+import org.jetbrains.org.objectweb.asm.tree.InsnNode
 import org.jetbrains.org.objectweb.asm.tree.MethodNode

 internal fun MethodNode.acceptWithStateMachine(
@@ -154,3 +160,30 @@ internal fun createFakeContinuation(context: JvmBackendContext): IrExpression =
    context.ir.symbols.continuationClass.createType(true, listOf(makeTypeProjection(context.irBuiltIns.anyNType, Variance.INVARIANT))),
    "FAKE_CONTINUATION"
 )
+
+fun MethodNode.preprocessSuspendMarkers(method: IrFunction? = null) {
+    if (instructions.first == null) return
+    if (method?.origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE) {
+        // Remove the fake continuation constructor, since this method either has a real state machine or is inline.
+        // Include one instruction before the start marker (that's the id) and one after the end marker (that's a pop).
+        val sequence = instructions.asSequence()
+        val start = sequence.firstOrNull { it.next != null && isBeforeFakeContinuationConstructorCallMarker(it.next) }
+        if (start != null) {
+            val end = sequence.first { it.previous != null && isAfterFakeContinuationConstructorCallMarker(it.previous) }
+            InsnSequence(start, end.next).forEach(instructions::remove)
+        }
+    }
+    // Method is null if this node is going to be inlined into another node rather than written into a class file.
+    val keepMarkers = method != null &&
+            method.origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE &&
+            method.origin != JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
+    for (insn in instructions.asSequence().filter { isBeforeInlineSuspendMarker(it) || isAfterInlineSuspendMarker(it) }) {
+        if (keepMarkers) {
+            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))
+        } else {
+            instructions.remove(insn.previous)
+            instructions.remove(insn)
+        }
+    }
+}
@@ -7,8 +7,6 @@ package org.jetbrains.kotlin.backend.jvm.codegen

 import org.jetbrains.kotlin.backend.common.lower.BOUND_RECEIVER_PARAMETER
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
-import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE
-import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
 import org.jetbrains.kotlin.backend.jvm.intrinsics.IrIntrinsicMethods
 import org.jetbrains.kotlin.backend.jvm.intrinsics.JavaClassProperty
 import org.jetbrains.kotlin.backend.jvm.lower.MultifileFacadeFileEntry
@@ -223,9 +221,7 @@ class ExpressionCodegen(
    }

    private fun generateFakeContinuationConstructorIfNeeded() {
-        if (irFunction.origin != FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE) return
-        val continuationClass = irFunction.suspendForInlineToOriginal()?.continuationClass()
-            ?: error("could not find continuation for ${irFunction.render()}")
+        val continuationClass = irFunction.suspendForInlineToOriginal()?.continuationClass() ?: return
        val continuationType = typeMapper.mapClass(continuationClass)
        val continuationIndex = frameMap.getIndex(irFunction.continuationParameter()!!.symbol)
        with(mv) {
@@ -378,6 +374,7 @@ class ExpressionCodegen(
        val callable = methodSignatureMapper.mapToCallableMethod(irFunction, expression)
        val callGenerator = getOrCreateCallGenerator(expression, data, callable.signature)
        val asmType = if (expression is IrConstructorCall) typeMapper.mapTypeAsDeclaration(expression.type) else expression.asmType
+        val isSuspensionPoint = expression.isSuspensionPoint()

        when {
            expression is IrConstructorCall -> {
@@ -406,7 +403,7 @@ class ExpressionCodegen(
            }
            expression.symbol.descriptor is ConstructorDescriptor ->
                throw AssertionError("IrCall with ConstructorDescriptor: ${expression.javaClass.simpleName}")
-            expression.isSuspensionPoint() ->
+            isSuspensionPoint != SuspensionPointKind.NEVER ->
                addInlineMarker(mv, isStartNotEnd = true)
        }

@@ -431,9 +428,8 @@ class ExpressionCodegen(

        expression.markLineNumber(true)

-        // Do not generate redundant markers in continuation class.
-        if (expression.isSuspensionPoint()) {
-            addSuspendMarker(mv, isStartNotEnd = true)
+        if (isSuspensionPoint != SuspensionPointKind.NEVER) {
+            addSuspendMarker(mv, isStartNotEnd = true, isInline = isSuspensionPoint == SuspensionPointKind.NOT_INLINE)
        }

        if (irFunction.isInvokeSuspendOfContinuation()) {
@@ -445,8 +441,8 @@ class ExpressionCodegen(
            callGenerator.genCall(callable, this, expression)
        }

-        if (expression.isSuspensionPoint()) {
-            addSuspendMarker(mv, isStartNotEnd = false)
+        if (isSuspensionPoint != SuspensionPointKind.NEVER) {
+            addSuspendMarker(mv, isStartNotEnd = false, isInline = isSuspensionPoint == SuspensionPointKind.NOT_INLINE)
            addInlineMarker(mv, isStartNotEnd = false)
        }

@@ -482,19 +478,18 @@ class ExpressionCodegen(
        }
    }

-    private fun IrFunctionAccessExpression.isSuspensionPoint(): Boolean = when {
-        !symbol.owner.isSuspend || !irFunction.shouldContainSuspendMarkers() -> false
+    private enum class SuspensionPointKind { NEVER, NOT_INLINE, ALWAYS }
+
+    private fun IrFunctionAccessExpression.isSuspensionPoint(): SuspensionPointKind = when {
+        !symbol.owner.isSuspend || !irFunction.shouldContainSuspendMarkers() -> SuspensionPointKind.NEVER
        // Copy-pasted bytecode blocks are not suspension points.
        symbol.owner.isInline ->
-            symbol.owner.let {
-                it.name.asString() == "suspendCoroutineUninterceptedOrReturn" &&
-                it.getPackageFragment()?.fqName == FqName("kotlin.coroutines.intrinsics")
-            }
+            if (symbol.owner.name.asString() == "suspendCoroutineUninterceptedOrReturn" &&
+                symbol.owner.getPackageFragment()?.fqName == FqName("kotlin.coroutines.intrinsics")
+            ) SuspensionPointKind.ALWAYS else SuspensionPointKind.NEVER
        // This includes inline lambdas, but only in functions intended for the inliner; in others, they stay as `f.invoke()`.
-        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
+        dispatchReceiver.isReadOfInlineLambda() -> SuspensionPointKind.NOT_INLINE
+        else -> SuspensionPointKind.ALWAYS
    }

    override fun visitVariable(declaration: IrVariable, data: BlockInfo): PromisedValue {
@@ -99,6 +99,7 @@ class IrSourceCompilerForInline(
    override fun generateLambdaBody(lambdaInfo: ExpressionLambda): SMAPAndMethodNode {
        val smap = codegen.context.getSourceMapper(codegen.classCodegen.irClass)
        val node = FunctionCodegen((lambdaInfo as IrExpressionLambdaImpl).function, codegen.classCodegen, codegen).generate(smap)
+        node.preprocessSuspendMarkers()
        return SMAPAndMethodNode(node, SMAP(smap.resultMappings))
    }

@@ -120,6 +121,7 @@ class IrSourceCompilerForInline(
            callee
        val classCodegen = FakeClassCodegen(forInlineFunction, codegen.classCodegen)
        val node = FunctionCodegen(forInlineFunction, classCodegen).generate()
+        node.preprocessSuspendMarkers()
        return SMAPAndMethodNode(node, SMAP(classCodegen.getOrCreateSourceMapper().resultMappings))
    }