Replace POP with ARETURN if it pops Unit and ARETURN shall return Unit

#KT-16880: Fixed
2017-11-10 15:46:11 +03:00
parent 2cdc246a27
commit c8904b1c7c
7 changed files with 354 additions and 2 deletions
@@ -2324,6 +2324,8 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
        callGenerator.genCall(callableMethod, resolvedCall, defaultMaskWasGenerated, this);

        if (isSuspendCall) {
+            addReturnsUnitMarkerIfNecessary(v, resolvedCall);
+
            addSuspendMarker(v, false);
            addInlineMarker(v, false);
        }
@@ -77,6 +77,8 @@ class CoroutineTransformerMethodVisitor(
        FixStackMethodTransformer().transform(containingClassInternalName, methodNode)

        if (isForNamedFunction) {
+            ReturnUnitMethodTransformer.transform(containingClassInternalName, methodNode)
+
            if (allSuspensionPointsAreTailCalls(containingClassInternalName, methodNode, suspensionPoints)) {
                dropSuspensionMarkers(methodNode, suspensionPoints)
                return
@@ -701,7 +703,7 @@ private fun allSuspensionPointsAreTailCalls(
    }
 }

-private class IgnoringCopyOperationSourceInterpreter : SourceInterpreter() {
+internal class IgnoringCopyOperationSourceInterpreter : SourceInterpreter() {
    override fun copyOperation(insn: AbstractInsnNode?, value: SourceValue?) = value
 }

@@ -0,0 +1,131 @@
+/*
+ * Copyright 2010-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.codegen.coroutines
+
+import org.jetbrains.kotlin.backend.common.pop
+import org.jetbrains.kotlin.codegen.inline.isReturnsUnitMarker
+import org.jetbrains.kotlin.codegen.optimization.boxing.isUnitInstance
+import org.jetbrains.kotlin.codegen.optimization.common.ControlFlowGraph
+import org.jetbrains.kotlin.codegen.optimization.common.asSequence
+import org.jetbrains.kotlin.codegen.optimization.common.isMeaningful
+import org.jetbrains.kotlin.codegen.optimization.common.removeAll
+import org.jetbrains.kotlin.codegen.optimization.transformer.MethodTransformer
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
+import org.jetbrains.kotlin.utils.keysToMap
+import org.jetbrains.kotlin.utils.sure
+import org.jetbrains.org.objectweb.asm.Opcodes
+import org.jetbrains.org.objectweb.asm.tree.*
+
+/*
+ * Replace POP with ARETURN iff
+ * 1) It is immediately followed by { GETSTATIC Unit.INSTANCE, ARETURN } sequences
+ * 2) It is poping Unit
+ */
+object ReturnUnitMethodTransformer : MethodTransformer() {
+    override fun transform(internalClassName: String, methodNode: MethodNode) {
+        val unitMarks = findReturnsUnitMarks(methodNode)
+        if (unitMarks.isEmpty()) return
+
+        val units = findReturnUnitSequences(methodNode)
+        if (units.isEmpty()) {
+            cleanUpReturnsUnitMarkers(methodNode, unitMarks)
+            return
+        }
+
+        val pops = methodNode.instructions.asSequence().filter { it.opcode == Opcodes.POP }.toList()
+        val popSuccessors = findSuccessors(methodNode, pops)
+        val sourceInsns = findSourceInstructions(internalClassName, methodNode, pops)
+        val safePops = filterOutUnsafes(popSuccessors, units, sourceInsns)
+
+        // Replace POP with ARETURN for tail call optimization
+        safePops.forEach { methodNode.instructions.set(it, InsnNode(Opcodes.ARETURN)) }
+        cleanUpReturnsUnitMarkers(methodNode, unitMarks)
+    }
+
+    // Return list of POPs, which can be safely replaced by ARETURNs
+    private fun filterOutUnsafes(
+            popSuccessors: Map<AbstractInsnNode, Collection<AbstractInsnNode>>,
+            units: Collection<AbstractInsnNode>,
+            sourceInsns: Map<AbstractInsnNode, Collection<AbstractInsnNode>>
+    ): Collection<AbstractInsnNode> {
+        return popSuccessors.filter { (pop, successors) ->
+            successors.all { it in units } &&
+            sourceInsns[pop].sure { "Sources of $pop cannot be null" }.all(::isSuspendingCallReturningUnit)
+        }.keys
+    }
+
+    // Find instructions which do something on stack, ignoring markers
+    // Return map {insn => list of found instructions}
+    private fun findSuccessors(
+            methodNode: MethodNode,
+            insns: List<AbstractInsnNode>
+    ): Map<AbstractInsnNode, Collection<AbstractInsnNode>> {
+        val cfg = ControlFlowGraph.build(methodNode)
+        return insns.keysToMap { findSuccessors(cfg, it, methodNode) }
+    }
+
+    // Find all meaningful successors of [insn]
+    private fun findSuccessors(cfg: ControlFlowGraph, insn: AbstractInsnNode, methodNode: MethodNode): Collection<AbstractInsnNode> {
+        val stack = cfg.getSuccessorsIndices(insn).mapTo(ArrayList()) { methodNode.instructions[it] }
+        val successors = arrayListOf<AbstractInsnNode>()
+        while (stack.isNotEmpty()) {
+            val current = stack.pop()
+            if (current in successors || isReturnsUnitMarker(current)) continue
+            if (!current.isMeaningful || current is JumpInsnNode || current.opcode == Opcodes.NOP) {
+                cfg.getSuccessorsIndices(current).mapTo(stack) { methodNode.instructions[it] }
+                continue
+            }
+            // There can be multiple chains of { UnitInstance, POP } after inlining. Ignore them
+            if (current.isUnitInstance()) {
+                val newSuccessors = findSuccessors(cfg, current, methodNode)
+                if (newSuccessors.all { it.opcode == Opcodes.POP }) {
+                    newSuccessors.flatMapTo(stack) { findSuccessors(cfg, it, methodNode) }
+                    continue
+                }
+            }
+            successors.add(current)
+        }
+        return successors
+    }
+
+    private fun isSuspendingCallReturningUnit(node: AbstractInsnNode): Boolean =
+            node.safeAs<MethodInsnNode>()?.next?.next?.let(::isReturnsUnitMarker) == true
+
+    private fun findSourceInstructions(
+            internalClassName: String,
+            methodNode: MethodNode,
+            pops: Collection<AbstractInsnNode>
+    ): Map<AbstractInsnNode, Collection<AbstractInsnNode>> {
+        val frames = analyze(internalClassName, methodNode, IgnoringCopyOperationSourceInterpreter())
+        return pops.keysToMap {
+            frames[methodNode.instructions.indexOf(it)].getStack(0).insns
+        }
+    }
+
+    // Find { GETSTATIC kotlin/Unit.INSTANCE, ARETURN } sequences
+    // Result is list of GETSTATIC kotlin/Unit.INSTANCE instructions
+    private fun findReturnUnitSequences(methodNode: MethodNode): Collection<AbstractInsnNode> =
+            methodNode.instructions.asSequence().filter { it.isUnitInstance() && it.next?.opcode == Opcodes.ARETURN }.toList()
+
+    private fun findReturnsUnitMarks(methodNode: MethodNode): Collection<AbstractInsnNode> =
+            methodNode.instructions.asSequence().filter(::isReturnsUnitMarker).toList()
+
+    private fun cleanUpReturnsUnitMarkers(methodNode: MethodNode, unitMarks: Collection<AbstractInsnNode>) {
+        unitMarks.forEach { methodNode.instructions.removeAll(listOf(it.previous, it)) }
+    }
+}
+
@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.codegen.inline

 import com.intellij.openapi.vfs.VirtualFile
+import org.jetbrains.kotlin.backend.common.descriptors.substitute
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.codegen.AsmUtil
 import org.jetbrains.kotlin.codegen.BaseExpressionCodegen
@@ -28,6 +29,7 @@ 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.unwrapInitialDescriptorForSuspendFunction
 import org.jetbrains.kotlin.codegen.intrinsics.classId
 import org.jetbrains.kotlin.codegen.optimization.common.intConstant
 import org.jetbrains.kotlin.codegen.state.GenerationState
@@ -43,14 +45,16 @@ import org.jetbrains.kotlin.load.kotlin.incremental.components.IncrementalCache
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes.ENUM_TYPE
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes.JAVA_CLASS_TYPE
 import org.jetbrains.kotlin.resolve.jvm.JvmClassName
 import org.jetbrains.kotlin.resolve.source.PsiSourceElement
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedCallableMemberDescriptor
-import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.util.OperatorNameConventions
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.org.objectweb.asm.*
 import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter
 import org.jetbrains.org.objectweb.asm.tree.*
@@ -87,6 +91,7 @@ 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
+private const val INLINE_MARKET_RETURNS_UNIT = 2
 private val INTRINSIC_ARRAY_CONSTRUCTOR_TYPE = AsmUtil.asmTypeByClassId(classId)

 internal fun getMethodNode(
@@ -386,6 +391,26 @@ internal fun addInlineMarker(v: InstructionAdapter, isStartNotEnd: Boolean) {
    )
 }

+internal fun addReturnsUnitMarkerIfNecessary(v: InstructionAdapter, resolvedCall: ResolvedCall<*>) {
+    val wrapperDescriptor = resolvedCall.candidateDescriptor.safeAs<FunctionDescriptor>() ?: return
+    val unsubstitutedDescriptor = wrapperDescriptor.unwrapInitialDescriptorForSuspendFunction()
+
+    val typeSubstitutor = TypeSubstitutor.create(
+            unsubstitutedDescriptor.typeParameters
+                    .withIndex()
+                    .associateBy({ it.value.typeConstructor }) {
+                        TypeProjectionImpl(resolvedCall.typeArguments[wrapperDescriptor.typeParameters[it.index]] ?: return)
+                    }
+    )
+
+    val substitutedDescriptor = unsubstitutedDescriptor.substitute(typeSubstitutor) ?: return
+    val returnType = substitutedDescriptor.returnType ?: return
+
+    if (KotlinBuiltIns.isUnit(returnType)) {
+        addReturnsUnitMarker(v)
+    }
+}
+
 internal fun addSuspendMarker(v: InstructionAdapter, isStartNotEnd: Boolean) {
    v.iconst(if (isStartNotEnd) INLINE_MARKER_BEFORE_SUSPEND_ID else INLINE_MARKER_AFTER_SUSPEND_ID)
    v.visitMethodInsn(
@@ -395,8 +420,18 @@ internal fun addSuspendMarker(v: InstructionAdapter, isStartNotEnd: Boolean) {
    )
 }

+private fun addReturnsUnitMarker(v: InstructionAdapter) {
+    v.iconst(INLINE_MARKET_RETURNS_UNIT)
+    v.visitMethodInsn(
+            Opcodes.INVOKESTATIC, INLINE_MARKER_CLASS_NAME,
+            "mark",
+            "(I)V", false
+    )
+}
+
 internal fun isBeforeSuspendMarker(insn: AbstractInsnNode) = isSuspendMarker(insn, INLINE_MARKER_BEFORE_SUSPEND_ID)
 internal fun isAfterSuspendMarker(insn: AbstractInsnNode) = isSuspendMarker(insn, INLINE_MARKER_AFTER_SUSPEND_ID)
+internal fun isReturnsUnitMarker(insn: AbstractInsnNode) = isSuspendMarker(insn, INLINE_MARKET_RETURNS_UNIT)

 private fun isSuspendMarker(insn: AbstractInsnNode, id: Int) =
        isInlineMarker(insn, "mark") && insn.previous.intConstant == id